#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""The section topics data pipeline is a sequence of :class:`pyspark.sql.DataFrame` extraction and transformation functions.
Inputs come from Wikimedia Foundation's `Analytics Data Lake <https://wikitech.wikimedia.org/wiki/Analytics/Data_Lake>`_:
- `Wikipedias wikitext <https://wikitech.wikimedia.org/wiki/Analytics/Data_Lake/Content/Mediawiki_wikitext_current>`_
(all Wikipedias by default as per
`wikipedias.txt <https://gitlab.wikimedia.org/repos/structured-data/image-suggestions-data-pipelines/-/blob/main/section_topics/section_topics/wikipedias.txt>`_
)
- `Wikidata item page links <https://wikitech.wikimedia.org/wiki/Analytics/Data_Lake/Content/Wikidata_item_page_link>`_
High-level steps:
- gather wikitext of sections at a given hierarchy level
via the `MediaWiki parser from hell <https://mwparserfromhell.readthedocs.io>`_.
Default: :const:`section_topics.pipeline.SECTION_LEVEL`, lead section included
- optionally filter out sections that don't convey relevant content, typically lists and tables
- extract Wikidata `QIDs <https://www.wikidata.org/wiki/Wikidata:Glossary#QID>`_
from `wikilinks <https://en.wikipedia.org/wiki/H:WIKILINK>`_:
the so-called **section topics**
- optionally filter out noisy topics, typically dates and numbers
- compute the relevance score
Output row example:
========== ======= ============== =========== ======== ======= ========== ============= ================= ========= =========== ===========
snapshot wiki_db page_namespace revision_id page_qid page_id page_title section_index section_title topic_qid topic_title topic_score
========== ======= ============== =========== ======== ======= ========== ============= ================= ========= =========== ===========
2023-01-16 enwiki 0 1127523670 Q36724 841 Attila 5 Solitary kingship Q3623581 Arnegisclus 1.13
========== ======= ============== =========== ======== ======= ========== ============= ================= ========= =========== ===========
More `documentation <https://www.mediawiki.org/wiki/Structured_Data_Across_Wikimedia/Section_Topics/Data_Pipeline>`_
lives in MediaWiki.
Functions are ordered by their execution in the pipeline.
"""
import argparse
import re
import string
from importlib import resources
from typing import List, Optional, Pattern, Tuple
import mwparserfromhell as mwp
from pyspark.sql import Column, DataFrame, SparkSession
from pyspark.sql import functions as F
from pyspark.sql import types as T
from section_topics import queries
# Spark settings for development on stat boxes.
# They should mirror production DAG properties:
# https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/blob/main/platform_eng/dags/section_topics_dag.py
DEV_SPARK_SETTINGS = {
'spark.driver.cores': 4,
'spark.driver.memory': '16g',
'spark.executor.cores': 4,
'spark.executor.memory': '8g',
'spark.dynamicAllocation.maxExecutors': 128,
'spark.executor.memoryOverhead': '4g',
'spark.sql.shuffle.partitions': 2048,
# Mitigate `FetchFailed` due to `Timeout`
'spark.shuffle.service.index.cache.size': 2048,
'spark.stage.maxConsecutiveAttempts': 10,
'spark.reducer.maxReqsInFlight': 1,
'spark.shuffle.io.maxRetries': 10,
'spark.shuffle.io.retryWait': '60s',
}
# `parse` UDF output schema
SECTIONS_SCHEMA = T.ArrayType(
T.StructType(
[
T.StructField('index', T.IntegerType(), nullable=False),
T.StructField('heading', T.StringType(), nullable=False),
# Empty array in case of no links or images
T.StructField('links', T.ArrayType(T.StringType()), nullable=False),
T.StructField('images', T.ArrayType(T.StringType()), nullable=False),
]
)
)
# Image extraction facilities
IMAGE_NAME_PREFIX = r"""
(?:
= # equals sign if in a template as named parameter
| \|(?!.+=) # or a pipe, as unnamed (positional) parameter in a template (i.e. not followed by `paramname=`)
| : # or a colon if of the form file:foo.jpg
| \n # or newline if in a gallery tag
)
"""
ALLOWED_IMAGE_NAME_CHARACTERS = r"""
# any character excluding carriage return, newline,
# hash, angle brackets, square brackets, vertical bar,
# colon, braces and forward slash
[^\r\n\#\<\>\[\]\|:\{\}/]
"""
# Combined regex to match image file names
IMAGE_RE = re.compile(
rf"""
{IMAGE_NAME_PREFIX} # prefix
({ALLOWED_IMAGE_NAME_CHARACTERS}+) # filename
\.(\w+) # extension with period
""",
re.VERBOSE,
)
# Based on https://commons.wikimedia.org/wiki/Special:MediaStatistics
IMAGE_EXTENSIONS = (
'bmp',
'djv',
'djvu',
'gif',
'jpg',
'jpeg',
'jpe',
'jps',
'png',
'apng',
'svg',
'tif',
'tiff',
'webp',
'xcf',
)
# Use these columns to apply the table filter.
# The default one has more columns, see `scripts/detect_html_tables.py`.
TABLE_FILTER_COLUMNS = ('wiki_db', 'page_id', 'section_title')
# Columns for denylist filter
SECTIONS_DENYLIST_COLUMNS = ('wiki_db', 'section_heading')
MEDIA_PREFIXES_PARQUET = 'section_topics/media_prefixes'
SECTIONS_DENYLIST_PARQUET = 'section_topics/section_titles_denylist'
# Documented constants
#: Section hierarchy level to be extracted.
#: Level 1 is just for page titles, actual sections start from level 2.
SECTION_LEVEL = 2
#: Reserved heading for the lead section, AKA section zero.
SECTION_ZERO_HEADING = '### zero ###'
#: ASCII punctuation characters to be stripped from section headings.
#: Include the ASCII white space, don't strip round brackets.
STRIP_CHARS = string.punctuation.replace('()', ' ')
#: All kinds of white space to be substituted for the ASCII one;
#: underscores turn into spaces as well.
SUBSTITUTE_PATTERN = r'[\s_]'
SUBSTITUTE_RE = re.compile(SUBSTITUTE_PATTERN)
#: Allowed HTML tags to be stripped from section headings.
#: Based on MediaWiki's
#: `Sanitizer <https://github.com/wikimedia/mediawiki/blob/e46398790d4f70c4288d4df0f38c609cd2e33808/includes/parser/Sanitizer.php#L153>`_
RECOGNIZED_HTML_TAGS = [
'tt',
'wbr',
'br',
'sup',
'var',
's',
'ruby',
'sub',
'ul',
'abbr',
'hr',
'p',
'li',
'ins',
'cite',
'pre',
'small',
'bdi',
'h6',
'em',
'h1',
'rp',
'table',
'strike',
'data',
'blockquote',
'i',
'div',
'mark',
'center',
'dd',
'tr',
'h4',
'u',
'rt',
'span',
'link',
'td',
'dt',
'th',
'strong',
'caption',
'big',
'rtc',
'dfn',
'time',
'code',
'dl',
'del',
'samp',
'b',
'font',
'rb',
'h5',
'bdo',
'ol',
'kbd',
'q',
'h2',
'meta',
'h3',
]
#: Match standard wikitext lists and tables:
#:
#: - line starts with ``*`` -> unordered list
#: - line starts with ``#`` -> ordered list + avoid link anchors
#: - ``{|`` anywhere -> table
LIST_OR_TABLE_RE = re.compile(r'^\*|^#|\{\|', re.MULTILINE)
#: Section content length threshold
MINIMUM_SECTION_CHARS = 500
[docs]
def load_articles(
spark: SparkSession,
wikis: str,
) -> DataFrame: # pragma: no cover
"""Load the current revision of all articles (redirects included) through the
:const:`section_topics.queries.ARTICLES` Data Lake query.
:param spark: an active Spark session
:param wikis: a string of comma-separated wikis to process
:return: the dataframe of all articles
"""
return spark.sql(queries.ARTICLES.format(wikis=wikis))
[docs]
def load_qids(
spark: SparkSession, weekly_snapshot: str, wikis: str
) -> DataFrame: # pragma: no cover
"""Load Wikidata QIDs with their page links through the
:const:`section_topics.queries.QIDS` Data Lake query.
:param spark: an active Spark session
:param weekly_snapshot: a ``YYYY-MM-DD`` date
:param wikis: a string of comma-separated wikis to process
:return: the dataframe of QIDs and page links
"""
return spark.sql(queries.QIDS.format(weekly=weekly_snapshot, wikis=wikis))
[docs]
def separate_redirects(articles_all: DataFrame) -> tuple:
"""Separate redirects from actual articles.
:param articles_all: a dataframe of all articles as returned by
:func:`load_articles`
:return: the dataframes of articles and redirects
"""
articles = articles_all.where(
(F.col('page_redirect_target').isNull()) | (F.col('page_redirect_target') == '')
)
redirects = articles_all.where(
(F.col('page_redirect_target').isNotNull())
& (F.col('page_redirect_target') != '')
)
return articles, redirects
[docs]
def apply_filter(
df: DataFrame, filter_df: DataFrame, broadcast: bool = False
) -> DataFrame:
"""Exclude rows of an input dataframe given a filter dataframe.
Anti-join all filter columns against input ones.
:param df: a dataframe to be filtered
:param filter_df: a dataframe acting as a filter. Columns must be a subset of ``df``
:param broadcast: whether to broadcast ``filter_df``, which tells Spark to perform a *broadcast hash join*,
i.e., :func:`pyspark.sql.functions.broadcast`. Much faster if ``filter_df`` is small
:return: the filtered ``df`` dataframe
"""
cols = set(df.columns)
filter_cols = set(filter_df.columns)
if not filter_cols.issubset(cols):
raise ValueError(
f'The filter dataframe has invalid columns: {filter_cols}. They must be a subset of {cols}'
)
if broadcast:
filter_df = F.broadcast(filter_df)
return df.join(filter_df, on=list(filter_cols), how='left_anti')
[docs]
def look_up_qids(pages: DataFrame, qids: DataFrame) -> DataFrame:
"""Look up page QIDs through page titles.
:param pages: a dataframe of pages as output by :func:`apply_filter`.
Pass the output of :func:`load_pages` if you want the full raw dataset.
:param qids: a dataframe of page IDs and Wikidata QIDs as output by :func:`load_qids`
:return: the ``pages`` dataframe with page QIDs added
"""
return pages.join(qids, on=['page_title', 'wiki_db'], how='inner').select(
'wiki_db',
'revision_id',
F.col('item_id').alias('page_qid'),
'page_id',
'page_title',
'revision_text',
)
[docs]
def wikitext_headings_to_anchors(headings: List[str]) -> List[str]:
"""Transform wikitext headings into URL anchors.
For instance, ``=== Album in studio ===`` becomes ``Album_in_studio``, and serves as a section link in
https://it.wikipedia.org/wiki/Gaznevada#Album_in_studio.
Anchors that occur more than once get a numeric suffix in the form ``anchor_N``.
:param headings: a list of wikitext headings
:return: the corresponding URL anchors
"""
def transform(text: str) -> str:
# Remove `=` and whitespace surrounding the heading text
text = re.sub(r'^(={1,6})\s*(.*?)\s*\1$', r'\2', text)
# Remove HTML tags
text = re.sub(fr'<\/?({"|".join(RECOGNIZED_HTML_TAGS)})(\s[^>]*)?>', '', text)
# Remove wikilink markup
text = re.sub(r'\[\[(.+?)\]\]', r'\1', text)
text = re.sub(r'(?<!\[)\[[^\[]*?(?=//).+? (.+?)\](?!\[)', r'\1', text)
# Remove italics (''), bold ('''), or combination (''''') markup
text = re.sub("'''''|'''|''", '', text)
# Spaces & underscores collapse to 1 underscore
text = re.sub('[ _]+', '_', text)
# Tabs to underscores, 1 for each tab
text = re.sub(r'\t', '_', text)
return text
anchors = [transform(heading) for heading in headings]
# Append the `_N` suffix to duplicate anchors
return [
a if not a in anchors[0:i] else f'{a}_{anchors[0:i].count(a) + 1}'
for i, a in enumerate(anchors)
]
[docs]
def normalize_heading_column(
column: str,
substitute_pattern: str = SUBSTITUTE_PATTERN,
strip_chars: str = STRIP_CHARS,
) -> Column:
"""Normalize a dataframe column of section headings for better matching.
Normalization steps:
- remove ``_N`` suffixes in case of duplicate section anchors
as added by :func:`wikitext_headings_to_anchors`
- replace characters matched by ``substitute_re`` with one ASCII white space
- strip leading and trailing ``strip_chars``
- lowercase
.. note::
This normalization is not perfect: it's a trade-off between several ones,
some of which may prevent from converging to a lowest common denominator.
However, only extreme edge cases might be affected.
:param column: a dataframe column name of section headings
:param substitute_pattern: (optional) a regular expression pattern
whose matches will be replaced by one ASCII white space
:param strip_chars: (optional) a string of characters to be stripped
:return: the column of normalized section headings
"""
# These chars will be surrounded by `[]` to form a regexp group, so escape literal `[` and `]`.
# Also double the backslash to prevent it from being an escape char itself.
escaped_strip_chars = (
strip_chars.replace('\\', '\\\\').replace('[', r'\[').replace(']', r'\]')
)
return F.lower(
F.regexp_replace(
F.regexp_replace(
F.regexp_replace(
column,
'_[0-9]+$',
'',
),
substitute_pattern,
' ',
),
f'^[{escaped_strip_chars}]*(.*?)[{escaped_strip_chars}]*$',
'$1',
)
)
def _has_list_or_table(
wikitext: Optional[str], list_or_table_re: Pattern = LIST_OR_TABLE_RE
) -> bool:
# Tell if the given wikitext contains at least one list or table.
if wikitext is not None and re.search(list_or_table_re, wikitext) is not None:
return True
return False
def _heading_contains_ref_tag(heading: str) -> bool:
# Tell if the given wikitext heading contains a reference tag
if re.search(r'<ref.+/>', heading) is not None:
return True
if (
re.search(r'<ref[^>]*>.*?</ref>', heading, flags=re.DOTALL | re.IGNORECASE)
is not None
):
return True
return False
def _is_content_longer_or_equal_than(
wikitext: Optional[str], minimum_section_size: int
) -> bool:
# Tell if the given wikitext is at least as long as minimum_section_size.
# If minimum_section_size <= 0, then we do no check, and return True. Useful for testing.
if minimum_section_size <= 0:
return True
if wikitext is not None and len(wikitext) >= minimum_section_size:
return True
return False
def _grab_images(wikitext: Optional[str]) -> List[str]:
# Grab all image file names from a wikitext
if wikitext is None:
return []
return [
f'{name}.{extension}'.strip().replace(' ', '_')
for name, extension in IMAGE_RE.findall(wikitext)
if extension.lower() in IMAGE_EXTENSIONS
]
[docs]
def parse_wikitext(keep_lists_and_tables: bool, minimum_section_size: int) -> F.udf:
"""`Currying <https://en.wikipedia.org/wiki/Currying>`_ function
that passes args to the underlying :func:`parse` PySpark user-defined function (UDF).
See also `this gist <https://gist.github.com/andrearota/5910b5c5ac65845f23856b2415474c38>`_.
:param keep_lists_and_tables: whether to keep sections with at least one standard wikitext list or table
:param minimum_section_size: minimum content character length for the section to be considered
:return: the actual UDF
"""
# NOTE `docs/api.rst` contains a similar docstring that goes into the public docs
@F.udf(returnType=SECTIONS_SCHEMA)
def parse(
wikitext: str,
section_level: int = SECTION_LEVEL,
section_zero_heading: str = SECTION_ZERO_HEADING,
) -> list:
"""Parse a wikitext string into a list of dicts.
Each dict represents a section within the given wikitext and has the following keys:
- ``index`` - section number, ``0`` for the leading section
- ``heading`` - section heading, ``section_zero_heading``'s value for the leading section
- ``links`` - list of section wikilinks
- ``images`` - list of section image file names
A section heading is normalized through :func:`wikitext_headings_to_anchors`.
Clean wikilinks with mwparserfromhell's
`strip_code() <https://mwparserfromhell.readthedocs.io/en/v0.6.4/api/mwparserfromhell.html#mwparserfromhell.wikicode.Wikicode.strip_code>`_.
This can lead to empty strings that should be filtered.
:param wikitext: a wikitext
:param section_level: a section hierarchy level to be extracted
:param section_zero_heading: a heading reserved to the lead section
:return: the list of *(index, heading, links, images)* section dictionaries
"""
parsed = []
parsed_wikicode = mwp.parse(wikitext, skip_style_tags=True)
sections = parsed_wikicode.get_sections(
levels=[section_level], include_lead=True, include_headings=True
)
headings = [str(heading) for heading in parsed_wikicode.filter_headings()]
anchors = wikitext_headings_to_anchors(headings)
heading_index = 0
for index, section in enumerate(sections):
try:
content = None
links = [
link.title.strip_code().replace(' ', '_')
for link in section.filter_wikilinks()
]
headings_in_section = section.filter_headings()
if len(headings_in_section) == 0:
heading = section_zero_heading
else:
heading = anchors[heading_index]
heading_index += len(headings_in_section)
section.remove(headings_in_section[0])
content = str(section)
if not _is_content_longer_or_equal_than(content, minimum_section_size):
continue
if not keep_lists_and_tables and _has_list_or_table(content):
continue
if _heading_contains_ref_tag(heading):
continue
parsed.append(
dict(
index=index,
heading=heading, # Case is preserved
links=links,
# Don't grab images from section zero, as `content` is `None`
images=_grab_images(content),
)
)
except (
KeyError
): # When mwp is unable to convert an invalid HTML entity to a Unicode point
continue
return parsed
return parse
[docs]
def normalize_wikilinks(link_column: str) -> Column:
"""Lowercase the first character of wikilink target titles.
`The link target is case-sensitive except for the first character
<https://en.wikipedia.org/wiki/Help:Link#Wikilinks_(internal_links)>`_.
:param link_column: a dataframe column name of wikilinks
:return: the column of normalized wikilinks. ``None`` values are kept.
"""
return F.expr(
f'replace(concat(lower(substring({link_column}, 1, 1)), substring({link_column}, 2)), " ", "_")'
)
[docs]
def clean_up_links(sections: DataFrame) -> DataFrame:
"""Filter empty strings and add a column of normalized wikilinks.
:param sections: a dataframe of sections and wikilinks as output by :func:`extract_sections`
:return: the cleaned ``sections`` dataframe
"""
return sections.where(F.col('link') != '').withColumn(
'normalized_link', normalize_wikilinks('link')
)
[docs]
def resolve_redirects(sections: DataFrame, redirects: DataFrame) -> DataFrame:
"""Follow section wikilinks redirects.
If a wikilink points to a redirect page, replace its title with the redirected one.
Normalize redirects via :func:`normalize_wikilinks`.
:param sections: a dataframe of sections and cleaned wikilinks as output by :func:`clean_up_links`
:param redirects: a dataframe of page titles and redirected page titles as output by :func:`load_redirects`
:return: the dataframe of redirected wikilinks. Both original and normalized ones are kept.
"""
return (
sections.alias('sections')
.join(
redirects.alias('redirects'),
on=[
F.col('sections.wiki_db') == F.col('redirects.wiki_db'),
F.col('sections.normalized_link')
== normalize_wikilinks('redirects.page_title'),
],
how='left',
)
.select(
'sections.wiki_db',
'sections.revision_id',
'page_qid',
'sections.page_id',
'sections.page_title',
'section_index',
'section_title',
# Keep both original and normalized links
F.coalesce('redirects.page_redirect_target', 'sections.link').alias('link'),
F.coalesce(
normalize_wikilinks('redirects.page_redirect_target'),
'sections.normalized_link',
).alias('normalized_link'),
)
)
[docs]
def gather_section_topics(sections: DataFrame, articles: DataFrame) -> DataFrame:
"""Align section wikilinks to their Wikidata QIDs: the so-called **section topics**.
:param sections: a dataframe of sections and normalized wikilinks as output by :func:`resolve_redirects`.
Pass the output of :func:`clean_up_links` to skip wikilinks pointing to redirect pages.
:param articles: a dataframe of articles as output by :func:`look_up_qids`
:return: the dataframe of sections and topics (as Wikidata QIDs). Original topic titles are kept.
"""
articles = articles.select('wiki_db', 'page_qid', 'page_title')
# Align links to their QIDs by joining on page titles.
# Maximize matches via normalization.
return (
sections.alias('sections')
.join(
articles.alias('articles'),
on=[
F.col('sections.wiki_db') == F.col('articles.wiki_db'),
F.col('sections.normalized_link')
== normalize_wikilinks('articles.page_title'),
],
how='left',
)
.select(
'sections.wiki_db',
'revision_id',
'sections.page_qid',
'page_id',
'sections.page_title',
'sections.section_index',
'sections.section_title',
F.concat_ws(
'|', 'sections.wiki_db', 'sections.page_qid', 'sections.section_index'
).alias('section_id'),
F.col('articles.page_qid').alias('topic_qid'),
F.coalesce(
'articles.page_title',
# TODO: some wikis actually allow non-capitalized titles,
# in which case below transformation may not be valid;
# see https://github.com/wikimedia/operations-mediawiki-config/blob/c67a67252fc70e81931c3af5d1a60902acff0438/wmf-config/InitialiseSettings.php#L65
F.concat(
F.upper(F.expr('substring(sections.link, 1, 1)')),
F.expr('substring(sections.link, 2)'),
),
).alias('topic_title'),
)
)
def _compute_tf(topics: DataFrame, base_column: str) -> DataFrame:
# Compute the relative term frequency (TF).
# TF is *within* a wiki if based on sections or *across* wikis if based on page QIDs.
# TF numerator: occurrences of *one* topic QID per section or page QID
tf_numerator = (
topics.groupBy(base_column, 'topic_qid')
.count()
.withColumnRenamed('count', 'tf_num')
)
# TF denominator: occurrences of *all* topic QIDs per section or page QID
tf_denominator = (
topics.groupBy(base_column).count().withColumnRenamed('count', 'tf_den')
)
# TF is the relative frequency
tf = tf_numerator.join(tf_denominator, on=[base_column]).withColumn(
'tf', F.col('tf_num') / F.col('tf_den')
)
# TF depends on the section/page and on the topic, so join on them
# NOTE Hitting https://issues.apache.org/jira/browse/SPARK-10925.
# Fix: join the other way around,
# see https://issues.apache.org/jira/browse/SPARK-14948?focusedCommentId=15432253#comment-15432253.
# The unit test doesn't cause this, though.
return tf.join(topics, on=[base_column, 'topic_qid'])
def _compute_idf(topics: DataFrame, base_column: str) -> DataFrame:
# Compute the intra-wiki inverse document frequency
# IDF numerator: count of section or page QID occurrences per wiki
idf_numerator = topics.groupBy('wiki_db').agg(
F.countDistinct(base_column).alias('idf_num')
)
# IDF denominator: count of sections or page QIDs where a topic QID occurs, per topic QID, per wiki
idf_denominator = topics.groupBy('wiki_db', 'topic_qid').agg(
F.countDistinct(base_column).alias('idf_den')
)
# IDF: log(numerator / denominator)
idf = idf_numerator.join(idf_denominator, on=['wiki_db']).withColumn(
'idf', F.log(F.col('idf_num') / F.col('idf_den'))
)
# IDF depends on the wiki and on the topic, so join on them
# NOTE Hitting https://issues.apache.org/jira/browse/SPARK-10925 again.
# `topic_qid` triggers the bug, `wiki_db` doesn't.
# Join the other way around again, plus alias both dataframes.
return idf.alias('idf').join(
topics.alias('topics'),
on=['wiki_db', 'topic_qid'],
)
[docs]
def compute_relevance(topics: DataFrame, level: str = 'section') -> DataFrame:
"""Compute either the section-level or the article-level relevance score for every section topic.
The section-level score is a standard
`term frequency-inverted document frequency <https://en.wikipedia.org/wiki/Tf%E2%80%93idf>`_
(TF-IDF).
The article-level score is a custom TF-IDF, where TF is **across wikis** and IDF is **within one wiki**.
Workflow:
1. filter null topic QIDs
2. compute TF numerator: occurrences of *one* topic QID in a section or page QID
3. compute TF denominator: occurrences of *all* topic QIDs in a section or page QID
4. compute TF: numerator / denominator
5. join with input on section or page QID and topic QID
6. compute IDF numerator: count of sections or page QIDs in a wiki
7. compute IDF denominator: count of sections or page QIDs where a topic QID occurs, in a wiki
8. compute IDF: log( numerator / denominator )
9. join with input on wiki and topic QID
10. compute TF-IDF: TF * IDF
:param topics: a dataframe of section topics as output by :func:`gather_section_topics`
:param level: (optional) at which level relevance is computed, ``section`` or ``article``
:return: the input dataframe with the ``tf_idf`` column added
"""
if level == 'section':
base_column = 'section_id'
elif level == 'article':
base_column = 'page_qid'
else:
raise ValueError(f"Invalid level: {level}. Must be 'section' or 'article'")
# Filter null topic QIDs
topics = topics.where(~F.isnull('topic_qid'))
topics = _compute_tf(topics, base_column)
topics = _compute_idf(topics, base_column)
# Compute TF-IDF
return topics.withColumn(
'tf_idf',
F.col('tf') * F.col('idf'),
)
[docs]
def compose_output(
scored_topics: DataFrame,
all_topics: DataFrame,
snapshot: str,
page_namespace: int = 0,
) -> DataFrame:
"""Fuse scored topics with null ones and build the output dataset.
:param scored_topics: a dataframe of scored topics as output by :func:`compute_relevance`
:param all_topics: a dataframe of all topics as output by :func:`gather_section_topics`
:param snapshot: a weekly snapshot to serve as the constant value for the ``snapshot`` column of the output dataframe
:param page_namespace: (optional) a page namespace to serve as the constant value for the ``page_namespace`` column
of the output dataframe
:return: the final output dataframe
"""
return scored_topics.unionByName(
all_topics.where(F.isnull('topic_qid')), allowMissingColumns=True
).select(
F.lit(snapshot).alias('snapshot'),
'wiki_db',
F.lit(page_namespace).alias('page_namespace'),
'revision_id',
'page_qid',
'page_id',
'page_title',
'section_index',
'section_title',
'topic_qid',
'topic_title',
F.col('tf_idf').alias('topic_score'),
)
def parse_args() -> argparse.Namespace: # pragma: no cover
parser = argparse.ArgumentParser(description='Gather section topics from Wikitext')
# Required.
# See current invocation in
# https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/blob/main/platform_eng/dags/section_topics_dag.py
parser.add_argument('snapshot', metavar='YYYY-MM-DD', help='snapshot date')
parser.add_argument(
'-o',
'--output',
metavar='hdfs_path/to/parquet',
required=True,
help='HDFS path to section topics output parquet',
)
parser.add_argument(
'-i',
'--images-output',
metavar='hdfs_path/to/parquet',
required=True,
help='HDFS path to section images output parquet',
)
parser.add_argument(
'-p',
'--page-filter',
metavar='hdfs_path/to/parquet',
required=True,
help=(
'HDFS path to parquet of (wiki, page revision ID) rows to exclude, '
'as output by "scripts/check_bad_parsing.py"'
),
)
parser.add_argument(
'-s',
'--section-title-filter',
metavar='hdfs_path/to/parquet',
required=True,
help=(
'HDFS path to parquet with a dataframe to exclude, '
'as output by "scripts/gather_section_titles_denylist.py". '
f'The dataframe must include {SECTIONS_DENYLIST_COLUMNS} columns'
),
)
parser.add_argument(
'-t',
'--table-filter',
metavar='hdfs_path/to/parquet',
required=True,
help=(
'HDFS path to parquet with a dataframe to exclude, '
'as output by "scripts/detect_html_tables.py". '
f'The dataframe must include {TABLE_FILTER_COLUMNS} columns'
),
)
parser.add_argument(
'-q',
'--qid-filter',
metavar='hdfs_path/to/parquet ...',
action='append',
required=True,
help=(
'HDFS path to parquet with a dataframe of Wikidata IDs to exclude, '
'as output by "scripts/fetch_qids_from_wikidata.py"'
),
)
# Optional
parser.add_argument(
'-w',
'--wikis',
metavar='/path/to/file.txt',
type=argparse.FileType('r'),
default=resources.open_text('section_topics', 'wikipedias.txt'),
help=(
'plain text file of wikis to process, one per line. '
'Default: all Wikipedias, see "wikipedias.txt"'
),
)
parser.add_argument(
'-l',
'--length-filter',
metavar='N',
type=int,
default=MINIMUM_SECTION_CHARS,
help=(
'exclude sections whose content length is less than the given number of characters. '
f'Default: {MINIMUM_SECTION_CHARS}'
),
)
parser.add_argument(
'--handle-media',
metavar='hdfs_path/to/parquet',
help=(
f'separate media links and dump them to the given path. '
'WARNING: the pipeline execution time will increase to roughly 20 minutes'
),
)
parser.add_argument(
'-m',
'--media-prefixes',
metavar='hdfs_path/to/parquet',
default=MEDIA_PREFIXES_PARQUET,
help=(
'HDFS path to parquet with a dataframe of media prefixes, '
'as output by "scripts/fetch_qids_from_wikidata.py". '
f'Default: see {MEDIA_PREFIXES_PARQUET}. '
'Ignored if "--handle-media" is not passed'
),
)
# Flags
parser.add_argument(
'--keep-lists-and-tables',
action='store_true',
help="don't skip sections with at least one standard wikitext list or table",
)
return parser.parse_args()
def main(args: argparse.Namespace, spark: SparkSession) -> None: # pragma: no cover
# Order CLI args by pipeline execution
weekly = args.snapshot
with args.wikis as win:
wikis = ', '.join([f"'{w.rstrip()}'" for w in win])
keep_lists_and_tables = args.keep_lists_and_tables
denylist = spark.read.parquet(args.section_title_filter)
exclusion_qids = spark.createDataFrame(
[], T.StructType([T.StructField('topic_qid', T.StringType(), nullable=False)])
)
for qid_filter_path in args.qid_filter:
qid_filter_df = spark.read.parquet(qid_filter_path)
exclusion_qids = exclusion_qids.union(qid_filter_df).distinct()
# Load inputs
articles_all = load_articles(spark, wikis)
qids = load_qids(spark, weekly, wikis)
articles_raw, redirects = separate_redirects(articles_all)
# Page filter
articles_bad = spark.read.parquet(args.page_filter)
articles_good = apply_filter(articles_raw, articles_bad)
# From articles to sections
articles = look_up_qids(articles_good, qids)
sections_with_links, sections_with_images = extract_sections(
articles, keep_lists_and_tables, args.length_filter
)
sections_with_images.write.parquet(args.images_output, mode='overwrite')
# Denylist filter
# (this is similar to apply_filter invocations below, but needs a special join
# to ensure proper normalization between section titles in both dataframes)
sections_with_links = sections_with_links.join(
denylist,
on=[
denylist.wiki_db == sections_with_links.wiki_db,
# normalize title to ensure a perfect string match against the denylist
normalize_heading_column(denylist.section_heading)
== normalize_heading_column(sections_with_links.section_title),
],
how='left_anti',
)
# Table filter
if not keep_lists_and_tables:
table_sections = spark.read.parquet(args.table_filter)
sections_with_links = apply_filter(
sections_with_links, table_sections.select(*TABLE_FILTER_COLUMNS)
)
# Media & wikilinks handling
if args.handle_media:
media_prefixes_df = spark.read.parquet(args.media_prefixes)
media_prefixes = [row.prefix for row in media_prefixes_df.collect()]
media, main_dataset = handle_media(sections_with_links, media_prefixes)
media.write.parquet(args.handle_media, mode='overwrite')
clean_sections = clean_up_links(main_dataset)
else:
clean_sections = clean_up_links(sections_with_links)
# Section topics
sections_with_links = resolve_redirects(clean_sections, redirects)
topics = gather_section_topics(sections_with_links, articles)
# QID filter
filtered = apply_filter(topics, exclusion_qids, broadcast=True)
# Uncomment to write the unscored dataset
# filtered.write.parquet(os.path.join(wd, weekly), mode='overwrite')
# Topic score
scored = compute_relevance(filtered)
output = compose_output(scored, filtered, weekly)
output.write.parquet(args.output, mode='overwrite')
if __name__ == '__main__':
args = parse_args()
# Uncomment for prod
spark = SparkSession.builder.getOrCreate()
# Uncomment for dev
# from wmfdata.spark import create_session
# spark = create_session(
# app_name='section-topics-dev',
# ship_python_env=True,
# extra_settings=DEV_SPARK_SETTINGS,
# )
main(args, spark)
spark.stop()