diff --git a/src/wikiq/__init__.py b/src/wikiq/__init__.py index c3d01cb..d9d9cc2 100755 --- a/src/wikiq/__init__.py +++ b/src/wikiq/__init__.py @@ -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, ) diff --git a/src/wikiq/resume.py b/src/wikiq/resume.py index dc3513c..1b3f602 100644 --- a/src/wikiq/resume.py +++ b/src/wikiq/resume.py @@ -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 diff --git a/test/Wikiq_Unit_Test.py b/test/Wikiq_Unit_Test.py index 4c22eb3..34b660a 100644 --- a/test/Wikiq_Unit_Test.py +++ b/test/Wikiq_Unit_Test.py @@ -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") - diff --git a/test/baseline_output/collapse-user_sailormoon.parquet b/test/baseline_output/collapse-user_sailormoon.parquet deleted file mode 100644 index 0cb1db1..0000000 Binary files a/test/baseline_output/collapse-user_sailormoon.parquet and /dev/null differ diff --git a/test/baseline_output/diff_sailormoon.parquet b/test/baseline_output/diff_sailormoon.parquet deleted file mode 100644 index 1b14699..0000000 Binary files a/test/baseline_output/diff_sailormoon.parquet and /dev/null differ diff --git a/test/baseline_output/text_sailormoon.parquet b/test/baseline_output/text_sailormoon.parquet deleted file mode 100644 index a8037b1..0000000 Binary files a/test/baseline_output/text_sailormoon.parquet and /dev/null differ diff --git a/test/test_resume.py b/test/test_resume.py index 2684c7c..2928fba 100644 --- a/test/test_resume.py +++ b/test/test_resume.py @@ -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: diff --git a/test/wikiq_test_utils.py b/test/wikiq_test_utils.py index 77b0caf..337c67d 100644 --- a/test/wikiq_test_utils.py +++ b/test/wikiq_test_utils.py @@ -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)