Benjamin Mako Hill
015d4f9164
Nate has given up on getting wikiq to output parquet directly because it
makes too many unpredictable large memory allocations. The workflow now
outputs JSONL in a single pass and then uses spark (wikiq_spark) to index
it as parquet in a second pass.
Remove the parquet output path along with the machinery that existed only
to support it: checkpoint files, resume temp-file merging, namespace
partitioning, and file rotation (--partition-namespaces,
--max-revisions-per-file). Resume support remains for JSONL output, where
the resume point is derived from the last complete line of the output
file. Also remove the parquet tests and baseline files.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit 72a8851373)
79 lines
2.8 KiB
Python
79 lines
2.8 KiB
Python
import os
|
|
import shutil
|
|
import subprocess
|
|
from typing import Final, Union
|
|
|
|
TEST_DIR: Final[str] = os.path.dirname(os.path.realpath(__file__))
|
|
WIKIQ: Final[str] = os.path.join(os.path.join(TEST_DIR, ".."), "src/wikiq/__init__.py")
|
|
TEST_OUTPUT_DIR: Final[str] = os.path.join(TEST_DIR, "test_output")
|
|
BASELINE_DIR: Final[str] = os.path.join(TEST_DIR, "baseline_output")
|
|
|
|
IKWIKI: Final[str] = "ikwiki-20180301-pages-meta-history"
|
|
SAILORMOON: Final[str] = "sailormoon"
|
|
TWINPEAKS: Final[str] = "twinpeaks"
|
|
REGEXTEST: Final[str] = "regextest"
|
|
|
|
|
|
class WikiqTester:
|
|
def __init__(
|
|
self,
|
|
wiki: str,
|
|
case_name: str,
|
|
suffix: Union[str, None] = None,
|
|
in_compression: str = "bz2",
|
|
baseline_format: str = "tsv",
|
|
out_format: str = "tsv",
|
|
):
|
|
self.input_file = os.path.join(
|
|
TEST_DIR, "dumps", "{0}.xml.{1}".format(wiki, in_compression)
|
|
)
|
|
|
|
basename = "{0}_{1}".format(case_name, wiki)
|
|
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:
|
|
shutil.rmtree(self.output)
|
|
|
|
# For JSONL, self.output is a file path. Create parent directory if needed.
|
|
if out_format == "jsonl":
|
|
parent_dir = os.path.dirname(self.output)
|
|
if parent_dir:
|
|
os.makedirs(parent_dir, exist_ok=True)
|
|
|
|
if suffix is None:
|
|
self.wikiq_baseline_name = "{0}.{1}".format(wiki, baseline_format)
|
|
self.wikiq_out_name = "{0}.{1}".format(wiki, out_format)
|
|
else:
|
|
self.wikiq_baseline_name = "{0}_{1}.{2}".format(
|
|
wiki, suffix, baseline_format
|
|
)
|
|
self.wikiq_out_name = "{0}_{1}.{2}".format(wiki, suffix, out_format)
|
|
|
|
if case_name is not None:
|
|
self.baseline_file = os.path.join(
|
|
BASELINE_DIR, "{0}_{1}".format(case_name, self.wikiq_baseline_name)
|
|
)
|
|
|
|
def call_wikiq(self, *args: str, out: bool = True):
|
|
"""
|
|
Calls wikiq with the passed arguments on the input file relevant to the test.
|
|
:param args: The command line arguments to pass to wikiq.
|
|
:param out: Whether to pass an output argument to wikiq.
|
|
:return: The output of the wikiq call.
|
|
"""
|
|
if out:
|
|
call = " ".join([WIKIQ, self.input_file, "-o", self.output, "--batch-size", "10", *args])
|
|
else:
|
|
call = " ".join([WIKIQ, self.input_file, "--batch-size", "10", *args])
|
|
|
|
print(call)
|
|
return subprocess.check_output(call, stderr=subprocess.PIPE, shell=True)
|