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

remove parquet output support

Browse Source 
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>
This commit is contained in:
Benjamin Mako Hill
2026-08-06 15:31:48 -07:00
parent dbcae5c64e
commit 72a8851373
8 changed files with 39 additions and 875 deletions
Download Patch File Download Diff File Expand all files Collapse all files

264
src/wikiq/__init__.py
View File

@@ -30,14 +30,7 @@ import wikiq.tables as tables
from wikiq.tables import RevisionTable
from wikiq.wiki_diff_matcher import WikiDiffMatcher
from wikiq.wikitext_parser import WikitextParser
from wikiq.resume import (
get_checkpoint_path,
read_checkpoint,
get_resume_point,
setup_resume_temp_output,
finalize_resume_merge,
cleanup_interrupted_resume,
)
from wikiq.resume import get_resume_point
TO_ENCODE = ("title", "editor")
PERSISTENCE_RADIUS = 7
@@ -46,7 +39,6 @@ from pathlib import Path
import pyarrow as pa
import pyarrow.csv as pacsv
import pyarrow.parquet as pq
from deltas import SegmentMatcher, SequenceMatcher
@@ -518,26 +510,21 @@ class WikiqParser:
revert_radius: int = 15,
output_jsonl: bool = False,
output_jsonl_dir: bool = False,
output_parquet: bool = False,
batch_size: int = 1024,
resume_point: Union[tuple, dict, None] = None,
partition_namespaces: bool = False,
resume_point: Union[tuple, None] = None,
external_links: bool = False,
citations: bool = False,
wikilinks: bool = False,
templates: bool = False,
headings: bool = False,
time_limit_seconds: Union[float, None] = None,
max_revisions_per_file: int = 0,
input_filename: Union[str, None] = None,
):
"""
Parameters:
persist : what persistence method to use. Takes a PersistMethod value
resume_point : if set, either a (pageid, revid) tuple for single-file output,
or a dict mapping namespace -> (pageid, revid) for partitioned output.
For single-file: skip all revisions up to and including this point.
max_revisions_per_file : if > 0, close and rotate output files after this many revisions
resume_point : if set, a (pageid, revid) tuple; skip all revisions up
to and including this point.
input_filename : original input filename (needed for .jsonl.d output to derive output filename)
"""
self.input_file = input_file
@@ -549,7 +536,6 @@ class WikiqParser:
self.revert_radius = revert_radius
self.diff = diff
self.text = text
self.partition_namespaces = partition_namespaces
self.resume_point = resume_point
self.external_links = external_links
self.citations = citations
@@ -558,7 +544,6 @@ class WikiqParser:
self.headings = headings
self.shutdown_requested = False
self.time_limit_seconds = time_limit_seconds
self.max_revisions_per_file = max_revisions_per_file
if namespaces is not None:
self.namespace_filter = set(namespaces)
else:
@@ -576,13 +561,9 @@ class WikiqParser:
self.batch_size = batch_size
self.output_jsonl = output_jsonl
self.output_jsonl_dir = output_jsonl_dir
self.output_parquet = output_parquet
self.output_file = output_file
if output_parquet:
self.pq_writer = None
self.parquet_buffer = []
elif output_jsonl:
if output_jsonl:
pass # JSONLWriter created in process()
else:
# TSV output
@@ -592,10 +573,6 @@ class WikiqParser:
else:
self.output_file = open(output_file, "wb")
# Checkpoint for tracking resume point (path only, no open file handle for NFS safety)
self.checkpoint_path = None
self.checkpoint_state = {} # namespace -> (pageid, revid) or None -> (pageid, revid)
def request_shutdown(self):
"""Request graceful shutdown. The process() method will exit after completing the current batch."""
self.shutdown_requested = True
@@ -620,76 +597,6 @@ class WikiqParser:
if timer is not None:
timer.cancel()
def _get_part_path(self, base_path, part_num):
"""Generate path with part number inserted before extension.
Example: output.parquet -> output.part0.parquet
"""
path = Path(base_path)
return path.parent / f"{path.stem}.part{part_num}{path.suffix}"
def _open_checkpoint(self, output_file):
"""Enable checkpointing for Parquet output only.
JSONL doesn't need checkpoint files - resume point is derived from last line.
"""
if not self.output_parquet or output_file == sys.stdout.buffer:
return
self.checkpoint_path = get_checkpoint_path(output_file, self.partition_namespaces)
Path(self.checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
print(f"Checkpoint enabled: {self.checkpoint_path}", file=sys.stderr)
def _update_checkpoint(self, pageid, revid, namespace=None, part=0):
"""Update checkpoint state and write atomically (NFS-safe)."""
if self.checkpoint_path is None:
return
if self.partition_namespaces:
self.checkpoint_state[namespace] = {"pageid": pageid, "revid": revid, "part": part}
else:
self.checkpoint_state = {"pageid": pageid, "revid": revid, "part": part}
# Atomic write: write to temp file, then rename
temp_path = self.checkpoint_path + ".tmp"
with open(temp_path, 'w') as f:
json.dump(self.checkpoint_state, f)
os.replace(temp_path, self.checkpoint_path)
def _close_checkpoint(self, delete=False):
"""Clean up checkpoint, optionally deleting it."""
if self.checkpoint_path is None:
return
if delete and os.path.exists(self.checkpoint_path):
os.remove(self.checkpoint_path)
print(f"Checkpoint deleted: {self.checkpoint_path}", file=sys.stderr)
elif os.path.exists(self.checkpoint_path):
print(f"Checkpoint preserved for resume: {self.checkpoint_path}", file=sys.stderr)
# Clean up any leftover temp file
temp_path = self.checkpoint_path + ".tmp"
if os.path.exists(temp_path):
os.remove(temp_path)
def _write_batch(self, row_buffer, schema, writer, pq_writers, ns_base_paths, sorting_cols, namespace=None, part_numbers=None):
"""Write a batch of rows to the appropriate writer.
For partitioned output, creates writer lazily if needed.
Returns (writer, num_rows) - writer used and number of rows written.
"""
num_rows = len(row_buffer.get("revid", []))
if self.partition_namespaces and namespace is not None:
if namespace not in pq_writers:
base_path = ns_base_paths[namespace]
part_num = part_numbers.get(namespace, 0) if part_numbers else 0
if self.max_revisions_per_file > 0:
ns_path = self._get_part_path(base_path, part_num)
else:
ns_path = base_path
Path(ns_path).parent.mkdir(exist_ok=True, parents=True)
pq_writers[namespace] = pq.ParquetWriter(
ns_path, schema, flavor="spark", sorting_columns=sorting_cols
)
writer = pq_writers[namespace]
writer.write(pa.record_batch(row_buffer, schema=schema))
return writer, num_rows
def make_matchmake_pairs(self, patterns, labels) -> list[RegexPair]:
if (patterns is not None and labels is not None) and (
len(patterns) == len(labels)
@@ -747,26 +654,7 @@ class WikiqParser:
time_limit_timer = self._start_time_limit_timer()
# Track whether we've passed the resume point
if self.resume_point is None:
found_resume_point = True
elif self.partition_namespaces:
found_resume_point = {}
else:
found_resume_point = False
# When resuming with parquet, write new data to temp file/directory and merge at the end
original_output_file = None
temp_output_file = None
original_partition_dir = None
if self.resume_point is not None and self.output_parquet:
original_output_file, temp_output_file, original_partition_dir = \
setup_resume_temp_output(self.output_file, self.partition_namespaces)
if temp_output_file is not None:
self.output_file = temp_output_file
# Open checkpoint file for tracking resume point
checkpoint_output = original_output_file if original_output_file else self.output_file
self._open_checkpoint(checkpoint_output)
found_resume_point = self.resume_point is None
# Construct dump file iterator
dump = WikiqIterator(self.input_file, collapse_user=self.collapse_user)
@@ -804,48 +692,8 @@ class WikiqParser:
# Initialize writer
writer = None
sorting_cols = None
ns_base_paths = {}
pq_writers = {}
part_numbers = {}
if self.output_parquet:
pageid_sortingcol = pq.SortingColumn(schema.get_field_index("articleid"))
revid_sortingcol = pq.SortingColumn(schema.get_field_index("revid"))
sorting_cols = [pageid_sortingcol, revid_sortingcol]
if self.resume_point is not None:
if self.partition_namespaces:
for ns, resume_data in self.resume_point.items():
part_numbers[ns] = resume_data[2] if len(resume_data) > 2 else 0
else:
part_numbers[None] = self.resume_point[2] if len(self.resume_point) > 2 else 0
if not self.partition_namespaces:
if self.max_revisions_per_file > 0:
output_path_with_part = self._get_part_path(self.output_file, part_numbers.get(None, 0))
else:
output_path_with_part = self.output_file
writer = pq.ParquetWriter(
output_path_with_part,
schema,
flavor="spark",
sorting_columns=sorting_cols,
)
else:
output_path = Path(self.output_file)
if self.namespace_filter is not None:
namespaces = self.namespace_filter
else:
namespaces = self.namespaces.values()
ns_base_paths = {
ns: (output_path.parent / f"namespace={ns}") / output_path.name
for ns in namespaces
}
for ns in namespaces:
if ns not in part_numbers:
part_numbers[ns] = 0
elif self.output_jsonl:
if self.output_jsonl:
append_mode = self.resume_point is not None
if self.output_jsonl_dir:
# Create directory for JSONL output
@@ -891,34 +739,16 @@ class WikiqParser:
# Write buffer: accumulate rows before flushing
write_buffer = defaultdict(list)
buffer_count = 0
last_namespace = None
def flush_buffer():
nonlocal write_buffer, buffer_count, last_namespace
nonlocal write_buffer, buffer_count
if buffer_count == 0:
return
row_buffer = dict(write_buffer)
namespace = last_namespace
if self.output_parquet:
if self.partition_namespaces:
self._write_batch(
row_buffer, schema, writer, pq_writers, ns_base_paths,
sorting_cols, namespace=namespace, part_numbers=part_numbers
)
else:
writer.write(pa.record_batch(row_buffer, schema=schema))
elif self.output_jsonl:
if self.output_jsonl:
writer.write_batch(row_buffer)
else:
writer.write(pa.record_batch(row_buffer, schema=schema))
# Update checkpoint
last_pageid = row_buffer["articleid"][-1]
last_revid = row_buffer["revid"][-1]
part = part_numbers.get(namespace if self.partition_namespaces else None, 0)
self._update_checkpoint(last_pageid, last_revid,
namespace=namespace if self.partition_namespaces else None,
part=part)
write_buffer = defaultdict(list)
buffer_count = 0
@@ -1103,7 +933,6 @@ class WikiqParser:
for k, v in oldest_row.items():
write_buffer[k].append(v)
buffer_count += 1
last_namespace = page.mwpage.namespace
if buffer_count >= self.batch_size:
flush_buffer()
@@ -1120,7 +949,6 @@ class WikiqParser:
for k, v in row.items():
write_buffer[k].append(v)
buffer_count += 1
last_namespace = page.mwpage.namespace
if buffer_count >= self.batch_size:
flush_buffer()
@@ -1146,7 +974,6 @@ class WikiqParser:
for k, v in pending_row.items():
write_buffer[k].append(v)
buffer_count += 1
last_namespace = page.mwpage.namespace
if self.shutdown_requested:
break
@@ -1170,24 +997,9 @@ class WikiqParser:
)
# Close all writers
if self.output_parquet and self.partition_namespaces:
for pq_writer in pq_writers.values():
pq_writer.close()
elif writer is not None:
if writer is not None:
writer.close()
# Close checkpoint file; delete it only if we completed without interruption
self._close_checkpoint(delete=not self.shutdown_requested)
# Merge temp output with original for parquet resume
if original_output_file is not None and temp_output_file is not None:
finalize_resume_merge(
original_output_file,
temp_output_file,
self.partition_namespaces,
original_partition_dir
)
def match_archive_suffix(input_filename):
if re.match(r".*\.7z$", input_filename):
cmd = ["7za", "x", "-so", input_filename]
@@ -1215,14 +1027,12 @@ def get_output_filename(input_filename, output_format='tsv') -> str:
Args:
input_filename: Input dump file path
output_format: 'tsv', 'jsonl', or 'parquet'
output_format: 'tsv' or 'jsonl'
"""
output_filename = re.sub(r"\.(7z|gz|bz2)?$", "", input_filename)
output_filename = re.sub(r"\.xml", "", output_filename)
if output_format == 'jsonl':
output_filename = output_filename + ".jsonl"
elif output_format == 'parquet':
output_filename = output_filename + ".parquet"
else:
output_filename = output_filename + ".tsv"
return output_filename
@@ -1249,7 +1059,7 @@ def main():
dest="output",
type=str,
nargs=1,
help="Output file or directory. Format is detected from extension: .jsonl for JSONL, .parquet for Parquet, otherwise TSV.",
help="Output file or directory. Format is detected from extension: .jsonl for JSONL, otherwise TSV.",
)
parser.add_argument(
@@ -1415,7 +1225,7 @@ def main():
dest="batch_size",
default=1500,
type=int,
help="How many revisions to process in each batch. This ends up being the Parquet row group size",
help="How many revisions to process in each batch.",
)
parser.add_argument(
@@ -1433,22 +1243,6 @@ def main():
help="Time limit in hours before graceful shutdown. Set to 0 to disable (default).",
)
parser.add_argument(
"--partition-namespaces",
dest="partition_namespaces",
action="store_true",
default=False,
help="For Parquet output, partition output by namespace into separate files.",
)
parser.add_argument(
"--max-revisions-per-file",
dest="max_revisions_per_file",
type=int,
default=0,
help="For Parquet output, split output into multiple files after this many revisions. Set to 0 to disable (default).",
)
args = parser.parse_args()
# set persistence method
@@ -1508,13 +1302,11 @@ def main():
# Detect output format from extension
output_jsonl_dir = output.endswith(".jsonl.d")
output_jsonl = output.endswith(".jsonl") or output_jsonl_dir
output_parquet = output.endswith(".parquet")
partition_namespaces = args.partition_namespaces and output_parquet
if args.stdout:
output_file = sys.stdout.buffer
elif output_jsonl or output_parquet:
# Output is a JSONL or Parquet file path - use it directly
elif output_jsonl:
# Output is a JSONL file path - use it directly
output_file = output
elif os.path.isdir(output):
# Output is a directory - derive filename from input
@@ -1526,25 +1318,14 @@ def main():
# Handle resume functionality before opening input file
resume_point = None
if args.resume:
if (output_jsonl or output_parquet) and not args.stdout:
# Clean up any interrupted resume from previous run
if output_parquet:
cleanup_result = cleanup_interrupted_resume(output_file, partition_namespaces)
if cleanup_result == "start_fresh":
resume_point = None
else:
resume_point = get_resume_point(output_file, partition_namespaces)
else:
# JSONL: get resume point from last line of file (no checkpoint)
resume_point = get_resume_point(output_file, input_file=filename)
if output_jsonl and not args.stdout:
# JSONL: get resume point from last line of file
resume_point = get_resume_point(output_file, input_file=filename)
if resume_point is not None:
if isinstance(resume_point, dict):
print(f"Resuming from checkpoint for {len(resume_point)} namespaces", file=sys.stderr)
else:
pageid, revid = resume_point[0], resume_point[1]
print(f"Resuming from checkpoint: pageid={pageid}, revid={revid}", file=sys.stderr)
pageid, revid = resume_point[0], resume_point[1]
print(f"Resuming from: pageid={pageid}, revid={revid}", file=sys.stderr)
else:
sys.exit("Error: --resume only works with JSONL or Parquet output (not stdout or TSV)")
sys.exit("Error: --resume only works with JSONL output (not stdout or TSV)")
# Now open the input file
print("Processing file: %s" % filename, file=sys.stderr)
@@ -1567,8 +1348,6 @@ def main():
diff=args.diff,
output_jsonl=output_jsonl,
output_jsonl_dir=output_jsonl_dir,
output_parquet=output_parquet,
partition_namespaces=partition_namespaces,
batch_size=args.batch_size,
resume_point=resume_point,
external_links=args.external_links,
@@ -1577,7 +1356,6 @@ def main():
templates=args.templates,
headings=args.headings,
time_limit_seconds=time_limit_seconds,
max_revisions_per_file=args.max_revisions_per_file,
input_filename=filename,
)

470
src/wikiq/resume.py
View File

@@ -1,10 +1,8 @@
"""
Checkpoint and resume functionality for wikiq output.
Resume functionality for wikiq JSONL output.
This module handles:
- Finding resume points in existing output (JSONL or Parquet)
- Merging resumed data with existing output (for Parquet, streaming, memory-efficient)
- Checkpoint file management for fast resume point lookup
The resume point is derived from the last complete line of the existing
output file, so no separate checkpoint files are needed.
"""
import json
@@ -12,147 +10,6 @@ import os
import sys
from collections import deque
import pyarrow.parquet as pq
def get_checkpoint_path(output_file, partition_namespaces=False):
"""Get the path to the checkpoint file for a given output file.
For partitioned output, the checkpoint is placed outside the partition directory
to avoid pyarrow trying to read it as a parquet file. The filename includes
the output filename to keep it unique per input file (for parallel jobs).
"""
if partition_namespaces:
partition_dir = os.path.dirname(output_file)
output_filename = os.path.basename(output_file)
parent_dir = os.path.dirname(partition_dir)
return os.path.join(parent_dir, output_filename + ".checkpoint")
return str(output_file) + ".checkpoint"
def read_checkpoint(checkpoint_path, partition_namespaces=False):
"""
Read resume point from checkpoint file if it exists.
Checkpoint format:
Single file: {"pageid": 54, "revid": 325} or {"pageid": 54, "revid": 325, "part": 2}
Partitioned: {"0": {"pageid": 54, "revid": 325, "part": 1}, ...}
Returns:
For single files: A tuple (pageid, revid) or (pageid, revid, part), or None if not found.
For partitioned: A dict mapping namespace -> (pageid, revid, part), or None.
"""
if not os.path.exists(checkpoint_path):
return None
try:
with open(checkpoint_path, 'r') as f:
data = json.load(f)
if not data:
return None
# Single-file format: {"pageid": ..., "revid": ..., "part": ...}
if "pageid" in data and "revid" in data:
part = data.get("part", 0)
if part > 0:
return (data["pageid"], data["revid"], part)
return (data["pageid"], data["revid"])
# Partitioned format: {"0": {"pageid": ..., "revid": ..., "part": ...}, ...}
result = {}
for key, value in data.items():
part = value.get("part", 0)
result[int(key)] = (value["pageid"], value["revid"], part)
return result if result else None
except (json.JSONDecodeError, IOError, KeyError, TypeError) as e:
print(f"Warning: Could not read checkpoint file {checkpoint_path}: {e}", file=sys.stderr)
return None
def cleanup_interrupted_resume(output_file, partition_namespaces):
"""
Merge any leftover .resume_temp files from a previous interrupted run.
This should be called BEFORE get_resume_point() so the resume point
is calculated from the merged data.
Returns:
None - no temp files found or normal merge completed
"start_fresh" - both original and temp were corrupted and deleted
"""
import shutil
if partition_namespaces:
partition_dir = os.path.dirname(output_file)
output_filename = os.path.basename(output_file)
temp_suffix = ".resume_temp"
if not os.path.isdir(partition_dir):
return
has_old_temp_files = False
for ns_dir in os.listdir(partition_dir):
if ns_dir.startswith('namespace='):
temp_path = os.path.join(partition_dir, ns_dir, output_filename + temp_suffix)
if os.path.exists(temp_path):
has_old_temp_files = True
break
if has_old_temp_files:
print(f"Found leftover temp files in {partition_dir} from previous interrupted partitioned run, merging first...", file=sys.stderr)
had_corruption = merge_partitioned_namespaces(partition_dir, temp_suffix, output_filename)
has_valid_data = False
for ns_dir in os.listdir(partition_dir):
if ns_dir.startswith('namespace='):
ns_path = os.path.join(partition_dir, ns_dir)
parquet_files = [f for f in os.listdir(ns_path) if f.endswith('.parquet') and not f.endswith('.resume_temp')]
if parquet_files:
has_valid_data = True
break
if had_corruption and not has_valid_data:
checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
if os.path.exists(checkpoint_path):
os.remove(checkpoint_path)
print("All partitioned files were corrupted, will start fresh.", file=sys.stderr)
return "start_fresh"
print("Previous temp files merged successfully.", file=sys.stderr)
else:
temp_output_file = output_file + ".resume_temp"
if os.path.exists(temp_output_file) and not os.path.isdir(temp_output_file):
print(f"Found leftover temp file {temp_output_file} from previous interrupted run, merging first...", file=sys.stderr)
merged_path = output_file + ".merged"
merged = merge_parquet_files(output_file, temp_output_file, merged_path)
if merged == "original_only":
os.remove(temp_output_file)
elif merged == "temp_only":
if os.path.exists(output_file):
os.remove(output_file)
os.rename(temp_output_file, output_file)
print("Recovered from temp file (original was corrupted or missing).", file=sys.stderr)
elif merged == "both_invalid":
if os.path.exists(output_file):
os.remove(output_file)
if os.path.exists(temp_output_file):
os.remove(temp_output_file)
checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
if os.path.exists(checkpoint_path):
os.remove(checkpoint_path)
print("Both files were corrupted, will start fresh.", file=sys.stderr)
return "start_fresh"
elif merged == "merged":
os.remove(output_file)
os.rename(merged_path, output_file)
os.remove(temp_output_file)
print("Previous temp file merged successfully.", file=sys.stderr)
else:
os.remove(temp_output_file)
def get_jsonl_resume_point(output_file, input_file=None):
"""Get resume point from last complete line of JSONL file.
@@ -207,325 +64,18 @@ def get_jsonl_resume_point(output_file, input_file=None):
return None
def get_resume_point(output_file, partition_namespaces=False, input_file=None):
def get_resume_point(output_file, input_file=None):
"""
Find the resume point(s) from existing output.
For JSONL: reads last line of file (no checkpoint needed).
For Parquet: checks checkpoint file, falls back to scanning parquet.
Find the resume point from existing JSONL output.
Args:
output_file: Path to the output file.
partition_namespaces: Whether the output uses namespace partitioning.
input_file: Path to input file (needed for .jsonl.d directory output).
Returns:
For single files: A tuple (pageid, revid) or (pageid, revid, part), or None.
For partitioned: A dict mapping namespace -> (pageid, revid, part), or None.
A (pageid, revid) tuple, or None.
"""
# For JSONL, read resume point directly from last line (no checkpoint needed)
if output_file.endswith('.jsonl') or output_file.endswith('.jsonl.d'):
result = get_jsonl_resume_point(output_file, input_file)
if result:
print(f"Resume point found from JSONL: pageid={result[0]}, revid={result[1]}", file=sys.stderr)
return result
# For Parquet, use checkpoint file (fast)
checkpoint_path = get_checkpoint_path(output_file, partition_namespaces)
checkpoint_result = read_checkpoint(checkpoint_path, partition_namespaces)
if checkpoint_result is not None:
print(f"Resume point found in checkpoint file {checkpoint_path}", file=sys.stderr)
return checkpoint_result
# Fall back to scanning parquet (slow, for backwards compatibility)
print(f"No checkpoint file found at {checkpoint_path}, scanning parquet output...", file=sys.stderr)
try:
if partition_namespaces:
return _get_resume_point_partitioned(output_file)
else:
return _get_resume_point_single_file(output_file)
except Exception as e:
print(f"Error reading resume point from {output_file}: {e}", file=sys.stderr)
return None
def _get_last_row_resume_point(pq_path):
"""Get resume point by reading only the last row group of a parquet file."""
pf = pq.ParquetFile(pq_path)
if pf.metadata.num_row_groups == 0:
return None
last_rg_idx = pf.metadata.num_row_groups - 1
table = pf.read_row_group(last_rg_idx, columns=['articleid', 'revid'])
if table.num_rows == 0:
return None
max_pageid = table['articleid'][-1].as_py()
max_revid = table['revid'][-1].as_py()
return (max_pageid, max_revid, 0)
def _get_resume_point_partitioned(output_file):
"""Find per-namespace resume points from partitioned output."""
partition_dir = os.path.dirname(output_file)
output_filename = os.path.basename(output_file)
if not os.path.exists(partition_dir) or not os.path.isdir(partition_dir):
return None
namespace_dirs = [d for d in os.listdir(partition_dir) if d.startswith('namespace=')]
if not namespace_dirs:
return None
resume_points = {}
for ns_dir in namespace_dirs:
ns = int(ns_dir.split('=')[1])
pq_path = os.path.join(partition_dir, ns_dir, output_filename)
if not os.path.exists(pq_path):
continue
try:
result = _get_last_row_resume_point(pq_path)
if result is not None:
resume_points[ns] = result
except Exception as e:
print(f"Warning: Could not read {pq_path}: {e}", file=sys.stderr)
continue
return resume_points if resume_points else None
def _get_resume_point_single_file(output_file):
"""Find resume point from a single parquet file."""
if not os.path.exists(output_file):
return None
if os.path.isdir(output_file):
return None
return _get_last_row_resume_point(output_file)
def merge_parquet_files(original_path, temp_path, merged_path):
"""
Merge two parquet files by streaming row groups.
Returns:
"merged" - merged file was created from both sources
"original_only" - temp was invalid, keep original unchanged
"temp_only" - original was corrupted but temp is valid
"both_invalid" - both files invalid
False - both files were valid but empty
"""
original_valid = False
temp_valid = False
original_pq = None
temp_pq = None
try:
original_pq = pq.ParquetFile(original_path)
original_valid = True
except Exception as e:
print(f"Warning: Original file {original_path} is corrupted or invalid: {e}", file=sys.stderr)
try:
if not os.path.exists(temp_path):
print(f"Note: Temp file {temp_path} does not exist", file=sys.stderr)
else:
temp_pq = pq.ParquetFile(temp_path)
temp_valid = True
except Exception:
print(f"Note: No new data in temp file {temp_path}", file=sys.stderr)
if not original_valid and not temp_valid:
print(f"Both original and temp files are invalid, will start fresh", file=sys.stderr)
return "both_invalid"
if not original_valid and temp_valid:
print(f"Original file corrupted but temp file is valid, recovering from temp", file=sys.stderr)
return "temp_only"
if original_valid and not temp_valid:
return "original_only"
merged_writer = None
for i in range(original_pq.num_row_groups):
row_group = original_pq.read_row_group(i)
if merged_writer is None:
merged_writer = pq.ParquetWriter(
merged_path,
row_group.schema,
flavor="spark"
)
merged_writer.write_table(row_group)
for i in range(temp_pq.num_row_groups):
row_group = temp_pq.read_row_group(i)
if merged_writer is None:
merged_writer = pq.ParquetWriter(
merged_path,
row_group.schema,
flavor="spark"
)
merged_writer.write_table(row_group)
if merged_writer is not None:
merged_writer.close()
return "merged"
return False
def merge_partitioned_namespaces(partition_dir, temp_suffix, file_filter):
"""
Merge partitioned namespace directories after resume.
Returns:
True if at least one namespace has valid data after merge
False if all namespaces ended up with corrupted/deleted data
"""
namespace_dirs = [d for d in os.listdir(partition_dir) if d.startswith('namespace=')]
had_corruption = False
expected_temp = file_filter + temp_suffix
for ns_dir in namespace_dirs:
ns_path = os.path.join(partition_dir, ns_dir)
temp_path = os.path.join(ns_path, expected_temp)
if not os.path.exists(temp_path):
continue
original_file = file_filter
original_path = os.path.join(ns_path, original_file)
if os.path.exists(original_path):
merged_path = original_path + ".merged"
merged = merge_parquet_files(original_path, temp_path, merged_path)
if merged == "original_only":
if os.path.exists(temp_path):
os.remove(temp_path)
elif merged == "temp_only":
os.remove(original_path)
os.rename(temp_path, original_path)
elif merged == "both_invalid":
if os.path.exists(original_path):
os.remove(original_path)
if os.path.exists(temp_path):
os.remove(temp_path)
had_corruption = True
elif merged == "merged":
os.remove(original_path)
os.rename(merged_path, original_path)
if os.path.exists(temp_path):
os.remove(temp_path)
else:
if os.path.exists(original_path):
os.remove(original_path)
if os.path.exists(temp_path):
os.remove(temp_path)
else:
try:
pq.ParquetFile(temp_path)
os.rename(temp_path, original_path)
except Exception:
if os.path.exists(temp_path):
os.remove(temp_path)
had_corruption = True
return had_corruption
def finalize_resume_merge(
original_output_file,
temp_output_file,
partition_namespaces,
original_partition_dir
):
"""
Finalize the resume by merging temp output with original output.
"""
import shutil
print("Merging resumed data with existing output...", file=sys.stderr)
try:
if partition_namespaces and original_partition_dir is not None:
file_filter = os.path.basename(original_output_file)
merge_partitioned_namespaces(original_partition_dir, ".resume_temp", file_filter)
if os.path.exists(temp_output_file) and os.path.isdir(temp_output_file):
shutil.rmtree(temp_output_file)
else:
merged_output_file = original_output_file + ".merged"
merged = merge_parquet_files(original_output_file, temp_output_file, merged_output_file)
if merged == "original_only":
if os.path.exists(temp_output_file):
os.remove(temp_output_file)
elif merged == "temp_only":
os.remove(original_output_file)
os.rename(temp_output_file, original_output_file)
elif merged == "both_invalid":
os.remove(original_output_file)
if os.path.exists(temp_output_file):
os.remove(temp_output_file)
elif merged == "merged":
os.remove(original_output_file)
os.rename(merged_output_file, original_output_file)
if os.path.exists(temp_output_file):
os.remove(temp_output_file)
else:
os.remove(original_output_file)
if os.path.exists(temp_output_file):
os.remove(temp_output_file)
print("Merge complete.", file=sys.stderr)
except Exception as e:
print(f"Error merging resume data for {original_output_file}: {e}", file=sys.stderr)
print(f"New data saved in: {temp_output_file}", file=sys.stderr)
raise
def setup_resume_temp_output(output_file, partition_namespaces):
"""
Set up temp output for resume mode (Parquet only).
Returns:
Tuple of (original_output_file, temp_output_file, original_partition_dir)
or (None, None, None) if no existing output to resume from.
"""
import shutil
original_output_file = None
temp_output_file = None
original_partition_dir = None
if partition_namespaces:
partition_dir = os.path.dirname(output_file)
output_filename = os.path.basename(output_file)
output_exists = False
if os.path.isdir(partition_dir):
for d in os.listdir(partition_dir):
if d.startswith('namespace='):
if os.path.exists(os.path.join(partition_dir, d, output_filename)):
output_exists = True
break
if output_exists:
original_partition_dir = partition_dir
else:
output_exists = isinstance(output_file, str) and os.path.exists(output_file)
if output_exists:
original_output_file = output_file
temp_output_file = output_file + ".resume_temp"
if os.path.exists(temp_output_file):
if os.path.isdir(temp_output_file):
shutil.rmtree(temp_output_file)
else:
os.remove(temp_output_file)
if partition_namespaces:
os.makedirs(temp_output_file, exist_ok=True)
return original_output_file, temp_output_file, original_partition_dir
result = get_jsonl_resume_point(output_file, input_file)
if result:
print(f"Resume point found from JSONL: pageid={result[0]}, revid={result[1]}", file=sys.stderr)
return result

39
test/Wikiq_Unit_Test.py
View File

@@ -440,7 +440,7 @@ def test_external_links_only():
# Verify citations column does NOT exist
assert "citations" not in test.columns, "citations column should NOT exist when only --external-links is used"
# Verify column has list/array type (pandas reads parquet lists as numpy arrays)
# Verify column has list/array type (pandas reads pyarrow lists as numpy arrays)
assert test["external_links"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
"external_links should be a list/array type or None"
@@ -495,7 +495,7 @@ def test_citations_only():
# Verify external_links column does NOT exist
assert "external_links" not in test.columns, "external_links column should NOT exist when only --citations is used"
# Verify column has list/array type (pandas reads parquet lists as numpy arrays)
# Verify column has list/array type (pandas reads pyarrow lists as numpy arrays)
assert test["citations"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
"citations should be a list/array type or None"
@@ -545,7 +545,7 @@ def test_external_links_and_citations():
assert "external_links" in test.columns, "external_links column should exist"
assert "citations" in test.columns, "citations column should exist"
# Verify both columns have list/array types (pandas reads parquet lists as numpy arrays)
# Verify both columns have list/array types (pandas reads pyarrow lists as numpy arrays)
assert test["external_links"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
"external_links should be a list/array type or None"
assert test["citations"].apply(lambda x: x is None or hasattr(x, '__len__')).all(), \
@@ -753,37 +753,4 @@ def test_headings():
print(f"Headings test passed! {len(test)} rows processed")
def test_parquet_output():
"""Test that Parquet output format works correctly."""
tester = WikiqTester(SAILORMOON, "parquet_output", in_compression="7z", out_format="parquet")
try:
tester.call_wikiq("--fandom-2020")
except subprocess.CalledProcessError as exc:
pytest.fail(exc.stderr.decode("utf8"))
# Verify output file exists
assert os.path.exists(tester.output), f"Parquet output file should exist at {tester.output}"
# Read and verify content
test = pd.read_parquet(tester.output)
# Verify expected columns exist
assert "revid" in test.columns
assert "articleid" in test.columns
assert "title" in test.columns
assert "namespace" in test.columns
# Verify row count matches JSONL output
tester_jsonl = WikiqTester(SAILORMOON, "parquet_compare", in_compression="7z", out_format="jsonl")
try:
tester_jsonl.call_wikiq("--fandom-2020")
except subprocess.CalledProcessError as exc:
pytest.fail(exc.stderr.decode("utf8"))
test_jsonl = pd.read_json(tester_jsonl.output, lines=True)
assert len(test) == len(test_jsonl), f"Parquet and JSONL should have same row count: {len(test)} vs {len(test_jsonl)}"
print(f"Parquet output test passed! {len(test)} rows")

BIN
test/baseline_output/collapse-user_sailormoon.parquet
View File

Binary file not shown.

BIN
test/baseline_output/diff_sailormoon.parquet
View File

Binary file not shown.

BIN
test/baseline_output/text_sailormoon.parquet
View File

Binary file not shown.

128
test/test_resume.py
View File

@@ -9,10 +9,6 @@ import time
import pytest
from wikiq.resume import (
get_checkpoint_path,
read_checkpoint,
)
from wikiq_test_utils import (
SAILORMOON,
TEST_DIR,
@@ -33,7 +29,7 @@ def read_jsonl(filepath):
def test_resume():
"""Test that --resume properly resumes processing from the last checkpoint."""
"""Test that --resume properly resumes processing from the last line of output."""
import pandas as pd
from pandas.testing import assert_frame_equal
@@ -57,10 +53,6 @@ def test_resume():
for row in full_rows[:middle_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[middle_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -103,10 +95,6 @@ def test_resume_with_diff():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[resume_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--diff", "--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -170,10 +158,6 @@ def test_resume_simple():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[resume_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--resume")
except subprocess.CalledProcessError as exc:
@@ -186,39 +170,6 @@ def test_resume_simple():
assert_frame_equal(df_full, df_resumed)
def test_checkpoint_read():
"""Test that read_checkpoint correctly reads checkpoint files."""
with tempfile.TemporaryDirectory() as tmpdir:
checkpoint_path = os.path.join(tmpdir, "test.jsonl.checkpoint")
# Test reading valid checkpoint
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": 100, "revid": 200}, f)
result = read_checkpoint(checkpoint_path)
assert result == (100, 200), f"Expected (100, 200), got {result}"
# Test reading non-existent checkpoint
result = read_checkpoint(os.path.join(tmpdir, "nonexistent.checkpoint"))
assert result is None, f"Expected None for non-existent file, got {result}"
# Test reading empty checkpoint
empty_path = os.path.join(tmpdir, "empty.checkpoint")
with open(empty_path, 'w') as f:
f.write("{}")
result = read_checkpoint(empty_path)
assert result is None, f"Expected None for empty checkpoint, got {result}"
# Test reading corrupted checkpoint
corrupt_path = os.path.join(tmpdir, "corrupt.checkpoint")
with open(corrupt_path, 'w') as f:
f.write("not valid json")
result = read_checkpoint(corrupt_path)
assert result is None, f"Expected None for corrupted checkpoint, got {result}"
print("Checkpoint read test passed!")
def test_resume_with_interruption():
"""Test that resume works correctly after interruption."""
import pandas as pd
@@ -245,9 +196,6 @@ def test_resume_with_interruption():
# Clean up for interrupted run
if os.path.exists(output_file):
os.remove(output_file)
checkpoint_path = get_checkpoint_path(output_file)
if os.path.exists(checkpoint_path):
os.remove(checkpoint_path)
# Start wikiq and interrupt it
cmd_partial = [
@@ -301,48 +249,6 @@ def test_resume_with_interruption():
assert_frame_equal(df_full, df_resumed)
def test_resume_parquet():
"""Test that --resume works correctly with Parquet output format."""
import pandas as pd
from pandas.testing import assert_frame_equal
import pyarrow.parquet as pq
tester_full = WikiqTester(SAILORMOON, "resume_parquet_full", in_compression="7z", out_format="parquet")
try:
tester_full.call_wikiq("--fandom-2020")
except subprocess.CalledProcessError as exc:
pytest.fail(exc.stderr.decode("utf8"))
full_output_path = tester_full.output
full_table = pq.read_table(full_output_path)
# Use unsorted indices consistently - slice the table and get checkpoint from same position
resume_idx = len(full_table) // 3
resume_revid = int(full_table.column("revid")[resume_idx].as_py())
resume_pageid = int(full_table.column("articleid")[resume_idx].as_py())
tester_partial = WikiqTester(SAILORMOON, "resume_parquet_partial", in_compression="7z", out_format="parquet")
partial_output_path = tester_partial.output
# Write partial Parquet file using the SAME schema as the full file
partial_table = full_table.slice(0, resume_idx + 1)
pq.write_table(partial_table, partial_output_path)
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": resume_pageid, "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
pytest.fail(exc.stderr.decode("utf8"))
df_full = full_table.to_pandas()
df_resumed = pd.read_parquet(partial_output_path)
assert_frame_equal(df_full, df_resumed)
def test_resume_tsv_error():
"""Test that --resume with TSV output produces a proper error message."""
tester = WikiqTester(SAILORMOON, "resume_tsv_error", in_compression="7z", out_format="tsv")
@@ -352,7 +258,7 @@ def test_resume_tsv_error():
pytest.fail("Expected error for --resume with TSV output")
except subprocess.CalledProcessError as exc:
stderr = exc.stderr.decode("utf8")
assert "Error: --resume only works with JSONL or Parquet" in stderr, \
assert "Error: --resume only works with JSONL output" in stderr, \
f"Expected proper error message, got: {stderr}"
print("TSV resume error test passed!")
@@ -387,10 +293,6 @@ def test_resume_data_equivalence():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[resume_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -433,10 +335,6 @@ def test_resume_with_persistence():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[resume_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--persistence wikidiff2", "--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -487,10 +385,6 @@ def test_resume_corrupted_jsonl_last_line():
# Record file size before resume
size_before = os.path.getsize(corrupt_output_path)
# NO checkpoint file - JSONL resume works from last valid line in the file
checkpoint_path = get_checkpoint_path(corrupt_output_path)
assert not os.path.exists(checkpoint_path), "Test setup error: checkpoint should not exist"
# Resume should detect corrupted line, truncate it, then append new data
try:
tester_corrupt.call_wikiq("--fandom-2020", "--resume")
@@ -537,10 +431,6 @@ def test_resume_diff_persistence_combined():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": full_rows[resume_idx]["articleid"], "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--diff", "--persistence wikidiff2", "--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -588,7 +478,7 @@ def test_resume_mid_page():
resume_revid = int(resume_rev["revid"])
resume_pageid = int(resume_rev["articleid"])
# Find global index for checkpoint
# Find global index for the resume point
global_idx = df_full[df_full["revid"] == resume_revid].index[0]
tester_partial = WikiqTester(SAILORMOON, "resume_midpage_partial", in_compression="7z", out_format="jsonl")
@@ -600,10 +490,6 @@ def test_resume_mid_page():
for row in rows_to_write:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": resume_pageid, "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--diff", "--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -653,10 +539,6 @@ def test_resume_page_boundary():
for row in rows_to_write:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": resume_pageid, "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--diff", "--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:
@@ -790,10 +672,6 @@ def test_resume_revert_detection():
for row in full_rows[:resume_idx + 1]:
f.write(json.dumps(row) + "\n")
checkpoint_path = get_checkpoint_path(partial_output_path)
with open(checkpoint_path, 'w') as f:
json.dump({"pageid": resume_pageid, "revid": resume_revid}, f)
try:
tester_partial.call_wikiq("--fandom-2020", "--resume")
except subprocess.CalledProcessError as exc:

13
test/wikiq_test_utils.py
View File

@@ -42,17 +42,8 @@ class WikiqTester:
else:
shutil.rmtree(self.output)
# Also clean up resume-related files
for temp_suffix in [".resume_temp", ".checkpoint", ".merged"]:
temp_path = self.output + temp_suffix
if os.path.exists(temp_path):
if os.path.isfile(temp_path):
os.remove(temp_path)
else:
shutil.rmtree(temp_path)
# For JSONL and Parquet, self.output is a file path. Create parent directory if needed.
if out_format in ("jsonl", "parquet"):
# 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)