collective/wikiq
20
0
Fork 0
Code Issues Pull Requests 3 Activity

Refactor revision parsing logic to be columnar #1

Merged
beason merged 27 commits from test-parquet into parquet_support 2025-06-17 18:22:26 +00:00
Conversation 10 Commits 27 Files Changed 24 +69 -68
3 changed files with 69 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit f916af9836 - Show all commits

1
.gitignore vendored
View File

@@ -13,4 +13,3 @@ uv.lock
# Python build and test output # Python build and test output
__pycache__/ __pycache__/
/test/test_output/ /test/test_output/
/test/test_output.parquet/

86
test/Wikiq_Unit_Test.py
View File

@@ -1,3 +1,4 @@
import shutil
import unittest import unittest
import os import os
import subprocess import subprocess
@@ -15,7 +16,6 @@ from typing import Final
TEST_DIR: Final[str] = os.path.dirname(os.path.realpath(__file__)) TEST_DIR: Final[str] = os.path.dirname(os.path.realpath(__file__))
WIKIQ: Final[str] = os.path.join(os.path.dirname(TEST_DIR), "wikiq") WIKIQ: Final[str] = os.path.join(os.path.dirname(TEST_DIR), "wikiq")
TEST_OUTPUT_DIR: Final[str] = os.path.join(TEST_DIR, "test_output") TEST_OUTPUT_DIR: Final[str] = os.path.join(TEST_DIR, "test_output")
PARQUET_OUTPUT_DIR: Final[str] = os.path.join(TEST_DIR, "test_output.parquet")
BASELINE_DIR: Final[str] = os.path.join(TEST_DIR, "baseline_output") BASELINE_DIR: Final[str] = os.path.join(TEST_DIR, "baseline_output")
IKWIKI: Final[str] = "ikwiki-20180301-pages-meta-history" IKWIKI: Final[str] = "ikwiki-20180301-pages-meta-history"
@@ -30,8 +30,6 @@ def setup():
# Perform directory check and reset here as this is a one-time setup step as opposed to per-test setup. # Perform directory check and reset here as this is a one-time setup step as opposed to per-test setup.
if not os.path.exists(TEST_OUTPUT_DIR): if not os.path.exists(TEST_OUTPUT_DIR):
os.mkdir(TEST_OUTPUT_DIR) os.mkdir(TEST_OUTPUT_DIR)
if not os.path.exists(PARQUET_OUTPUT_DIR):
os.mkdir(PARQUET_OUTPUT_DIR)
# Always run setup, even if this is executed via "python -m unittest" rather # Always run setup, even if this is executed via "python -m unittest" rather
@@ -42,7 +40,7 @@ setup()
class WikiqTester: class WikiqTester:
def __init__(self, def __init__(self,
wiki: str, wiki: str,
case_name: str | None = None, case_name: str,
suffix: str | None = None, suffix: str | None = None,
in_compression: str = "bz2", in_compression: str = "bz2",
baseline_format: str = "tsv", baseline_format: str = "tsv",
@@ -50,10 +48,20 @@ class WikiqTester:
): ):
self.input_file = os.path.join(TEST_DIR, "dumps", "{0}.xml.{1}".format(wiki, in_compression)) self.input_file = os.path.join(TEST_DIR, "dumps", "{0}.xml.{1}".format(wiki, in_compression))
if out_format == "tsv": basename = "{0}_{1}".format(case_name, wiki)
self.output_dir = TEST_OUTPUT_DIR if suffix:
basename = "{0}_{1}".format(basename, suffix)
self.output = os.path.join(TEST_OUTPUT_DIR, "{0}.{1}".format(basename, out_format))
if os.path.exists(self.output):
if os.path.isfile(self.output):
os.remove(self.output)
else: else:
self.output_dir = "{0}.parquet".format(TEST_OUTPUT_DIR) shutil.rmtree(self.output)
if out_format == "parquet":
os.makedirs(self.output, exist_ok=True)
if suffix is None: if suffix is None:
self.wikiq_baseline_name = "{0}.{1}".format(wiki, baseline_format) self.wikiq_baseline_name = "{0}.{1}".format(wiki, baseline_format)
@@ -61,14 +69,10 @@ class WikiqTester:
else: else:
self.wikiq_baseline_name = "{0}_{1}.{2}".format(wiki, suffix, baseline_format) self.wikiq_baseline_name = "{0}_{1}.{2}".format(wiki, suffix, baseline_format)
self.wikiq_out_name = "{0}_{1}.{2}".format(wiki, suffix, out_format) self.wikiq_out_name = "{0}_{1}.{2}".format(wiki, suffix, out_format)
self.call_output = os.path.join(self.output_dir, "{0}.{1}".format(wiki, out_format))
# If case_name is unset, there are no relevant baseline or test files. # If case_name is unset, there are no relevant baseline or test files.
if case_name is not None: if case_name is not None:
self.baseline_file = os.path.join(BASELINE_DIR, "{0}_{1}".format(case_name, self.wikiq_baseline_name)) self.baseline_file = os.path.join(BASELINE_DIR, "{0}_{1}".format(case_name, self.wikiq_baseline_name))
self.test_file = os.path.join(self.output_dir, "{0}_{1}".format(case_name, self.wikiq_out_name))
if os.path.exists(self.test_file):
os.remove(self.test_file)
def call_wikiq(self, *args: str, out: bool = True): def call_wikiq(self, *args: str, out: bool = True):
""" """
@@ -78,7 +82,7 @@ class WikiqTester:
:return: The output of the wikiq call. :return: The output of the wikiq call.
""" """
if out: if out:
call = ' '.join([WIKIQ, self.input_file, "-o", self.output_dir, *args]) call = ' '.join([WIKIQ, self.input_file, "-o", self.output, *args])
else: else:
call = ' '.join([WIKIQ, self.input_file, *args]) call = ' '.join([WIKIQ, self.input_file, *args])
@@ -105,9 +109,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -119,10 +121,8 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file)
# as a test let's make sure that we get equal data frames # as a test let's make sure that we get equal data frames
test = pd.read_table(tester.test_file) test = pd.read_table(tester.output)
num_wrong_ns = sum(~ test.namespace.isin({0, 1})) num_wrong_ns = sum(~ test.namespace.isin({0, 1}))
self.assertEqual(num_wrong_ns, 0) self.assertEqual(num_wrong_ns, 0)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
@@ -136,10 +136,8 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file)
# as a test let's make sure that we get equal data frames # as a test let's make sure that we get equal data frames
test = pd.read_table(tester.test_file) test = pd.read_table(tester.output)
num_wrong_ns = sum(~ test.namespace.isin({0, 1})) num_wrong_ns = sum(~ test.namespace.isin({0, 1}))
self.assertEqual(num_wrong_ns, 0) self.assertEqual(num_wrong_ns, 0)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
@@ -153,10 +151,8 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file)
# as a test let's make sure that we get equal data frames # as a test let's make sure that we get equal data frames
test = pd.read_table(tester.test_file) test = pd.read_table(tester.output)
num_reverted = sum(i is None for i in test.revert) num_reverted = sum(i is None for i in test.revert)
self.assertEqual(num_reverted, 0) self.assertEqual(num_reverted, 0)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
@@ -170,8 +166,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -183,9 +178,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -197,9 +190,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -211,9 +202,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -225,9 +214,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -239,16 +226,14 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
test = test.reindex(columns=sorted(test.columns)) test = test.reindex(columns=sorted(test.columns))
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
def test_malformed_noargs(self): def test_malformed_noargs(self):
tester = WikiqTester(wiki=TWINPEAKS, in_compression="7z") tester = WikiqTester(wiki=TWINPEAKS, case_name="noargs", in_compression="7z")
want_exception = 'xml.etree.ElementTree.ParseError: no element found: line 1369, column 0' want_exception = 'xml.etree.ElementTree.ParseError: no element found: line 1369, column 0'
try: try:
@@ -267,14 +252,12 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file)
test = pd.read_table(StringIO(outs)) test = pd.read_table(StringIO(outs))
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
def test_bad_regex(self): def test_bad_regex(self):
tester = WikiqTester(wiki=REGEXTEST) tester = WikiqTester(wiki=REGEXTEST, case_name="bad_regex")
# sample arguments for checking that bad arguments get terminated / test_regex_arguments # sample arguments for checking that bad arguments get terminated / test_regex_arguments
bad_arguments_list = [ bad_arguments_list = [
@@ -315,9 +298,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -337,9 +318,7 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file) test = pd.read_table(tester.output)
test = pd.read_table(tester.test_file)
baseline = pd.read_table(tester.baseline_file) baseline = pd.read_table(tester.baseline_file)
assert_frame_equal(test, baseline, check_like=True) assert_frame_equal(test, baseline, check_like=True)
@@ -352,10 +331,8 @@ class WikiqTestCase(unittest.TestCase):
except subprocess.CalledProcessError as exc: except subprocess.CalledProcessError as exc:
self.fail(exc.stderr.decode("utf8")) self.fail(exc.stderr.decode("utf8"))
copyfile(tester.call_output, tester.test_file)
# as a test let's make sure that we get equal data frames # as a test let's make sure that we get equal data frames
test: DataFrame = pd.read_parquet(tester.test_file) test: DataFrame = pd.read_parquet(tester.output)
# test = test.drop(['reverteds'], axis=1) # test = test.drop(['reverteds'], axis=1)
baseline: DataFrame = pd.read_table(tester.baseline_file) baseline: DataFrame = pd.read_table(tester.baseline_file)
@@ -364,6 +341,7 @@ class WikiqTestCase(unittest.TestCase):
baseline['date_time'] = pd.to_datetime(baseline['date_time']) baseline['date_time'] = pd.to_datetime(baseline['date_time'])
# Split strings to the arrays of reverted IDs so they can be compared. # Split strings to the arrays of reverted IDs so they can be compared.
baseline['revert'] = baseline['revert'].replace(np.nan, None) baseline['revert'] = baseline['revert'].replace(np.nan, None)
baseline['reverteds'] = baseline['reverteds'].replace(np.nan, None)
# baseline['reverteds'] = [None if i is np.nan else [int(j) for j in str(i).split(",")] for i in baseline['reverteds']] # baseline['reverteds'] = [None if i is np.nan else [int(j) for j in str(i).split(",")] for i in baseline['reverteds']]
baseline['sha1'] = baseline['sha1'].replace(np.nan, None) baseline['sha1'] = baseline['sha1'].replace(np.nan, None)
baseline['editor'] = baseline['editor'].replace(np.nan, None) baseline['editor'] = baseline['editor'].replace(np.nan, None)

44
wikiq
View File

@@ -201,6 +201,27 @@ class RegexPair(object):
return rev_data return rev_data
def pa_schema() -> pa.Schema:
fields = [
pa.field("revid", pa.int64()),
pa.field("date_time", pa.timestamp('s')),
pa.field("articleid", pa.int64()),
pa.field("editorid", pa.int64(), nullable=True),
pa.field("title", pa.string()),
pa.field("namespace", pa.int32()),
pa.field("deleted", pa.bool_()),
pa.field("text_chars", pa.int32()),
pa.field("comment_chars", pa.int32()),
pa.field("revert", pa.bool_(), nullable=True),
# reverteds is a string which contains a comma-separated list of reverted revision ids.
pa.field("reverteds", pa.string(), nullable=True),
pa.field("sha1", pa.string()),
pa.field("minor", pa.bool_()),
pa.field("editor", pa.string()),
pa.field("anon", pa.bool_())
]
return pa.schema(fields)
""" """
We used to use a dictionary to collect fields for the output. We used to use a dictionary to collect fields for the output.
@@ -229,6 +250,7 @@ class Revision:
namespace: int namespace: int
deleted: bool deleted: bool
text_chars: int | None = None text_chars: int | None = None
comment_chars: int | None = None
revert: bool | None = None revert: bool | None = None
reverteds: str = None reverteds: str = None
sha1: str | None = None sha1: str | None = None
@@ -249,6 +271,7 @@ class Revision:
pa.field("namespace", pa.int32()), pa.field("namespace", pa.int32()),
pa.field("deleted", pa.bool_()), pa.field("deleted", pa.bool_()),
pa.field("text_chars", pa.int32()), pa.field("text_chars", pa.int32()),
# pa.field("comment_chars", pa.int32()),
pa.field("revert", pa.bool_(), nullable=True), pa.field("revert", pa.bool_(), nullable=True),
# reverteds is a string which contains a comma-separated list of reverted revision ids. # reverteds is a string which contains a comma-separated list of reverted revision ids.
pa.field("reverteds", pa.string(), nullable=True), pa.field("reverteds", pa.string(), nullable=True),
@@ -492,6 +515,7 @@ class WikiqParser:
state = persistence.State() state = persistence.State()
# Iterate through a page's revisions # Iterate through a page's revisions
prev_text_chars = 0
for revs in page: for revs in page:
revs = list(revs) revs = list(revs)
rev = revs[-1] rev = revs[-1]
@@ -525,6 +549,7 @@ class WikiqParser:
# TODO rev.bytes doesn't work.. looks like a bug # TODO rev.bytes doesn't work.. looks like a bug
rev_data.text_chars = len(rev.text) rev_data.text_chars = len(rev.text)
rev_data.comment_chars = sum(0 if r.comment is None else len(r.comment) for r in revs)
# generate revert data # generate revert data
if rev_detector is not None: if rev_detector is not None:
@@ -550,7 +575,6 @@ class WikiqParser:
# TODO missing: additions_size deletions_size # TODO missing: additions_size deletions_size
# if collapse user was on, let's run that # if collapse user was on, let's run that
if self.collapse_user:
rev_data.collapsed_revs = len(revs) rev_data.collapsed_revs = len(revs)
# get the # get the
@@ -652,7 +676,7 @@ def main():
parser.add_argument('dumpfiles', metavar="DUMPFILE", nargs="*", type=str, parser.add_argument('dumpfiles', metavar="DUMPFILE", nargs="*", type=str,
help="Filename of the compressed or uncompressed XML database dump. If absent, we'll look for content on stdin and output on stdout.") help="Filename of the compressed or uncompressed XML database dump. If absent, we'll look for content on stdin and output on stdout.")
parser.add_argument('-o', '--output-dir', metavar='DIR', dest='output_dir', type=str, nargs=1, parser.add_argument('-o', '--output', metavar='OUTPUT', dest='output', type=str, nargs=1,
help="Directory for output files. If it ends with .parquet output will be in parquet format.") help="Directory for output files. If it ends with .parquet output will be in parquet format.")
parser.add_argument('-s', '--stdout', dest="stdout", action="store_true", parser.add_argument('-s', '--stdout', dest="stdout", action="store_true",
@@ -714,27 +738,27 @@ def main():
namespaces = None namespaces = None
if len(args.dumpfiles) > 0: if len(args.dumpfiles) > 0:
output_parquet = False
for filename in args.dumpfiles: for filename in args.dumpfiles:
input_file = open_input_file(filename, args.fandom_2020) input_file = open_input_file(filename, args.fandom_2020)
# open directory for output # open directory for output
if args.output_dir: if args.output:
output_dir = args.output_dir[0] output = args.output[0]
else: else:
output_dir = "." output = "."
if output_dir.endswith(".parquet"): output_parquet = output.endswith(".parquet")
output_parquet = True
print("Processing file: %s" % filename, file=sys.stderr) print("Processing file: %s" % filename, file=sys.stderr)
if args.stdout: if args.stdout:
# Parquet libraries need a binary output, so just sys.stdout doesn't work. # Parquet libraries need a binary output, so just sys.stdout doesn't work.
output_file = sys.stdout.buffer output_file = sys.stdout.buffer
else: elif os.path.isdir(output) or output_parquet:
filename = os.path.join(output_dir, os.path.basename(filename)) filename = os.path.join(output, os.path.basename(filename))
output_file = get_output_filename(filename, parquet=output_parquet) output_file = get_output_filename(filename, parquet=output_parquet)
else:
output_file = output
wikiq = WikiqParser(input_file, wikiq = WikiqParser(input_file,
output_file, output_file,