550 lines
23 KiB
Python
550 lines
23 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.dataset as ds
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
from pandas.testing import assert_frame_equal
|
|
|
|
from wikiq.resume import (
|
|
cleanup_interrupted_resume,
|
|
get_checkpoint_path,
|
|
get_resume_point,
|
|
merge_parquet_files,
|
|
)
|
|
from wikiq_test_utils import (
|
|
SAILORMOON,
|
|
TEST_DIR,
|
|
TEST_OUTPUT_DIR,
|
|
WIKIQ,
|
|
WikiqTester,
|
|
)
|
|
|
|
|
|
def test_resume():
|
|
"""Test that --resume properly resumes processing from the last written revid."""
|
|
tester_full = WikiqTester(SAILORMOON, "resume_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 = os.path.join(tester_full.output, f"{SAILORMOON}.parquet")
|
|
full_table = pq.read_table(full_output_path)
|
|
|
|
middle_idx = len(full_table) // 2
|
|
resume_revid = full_table.column("revid")[middle_idx].as_py()
|
|
|
|
print(f"Total revisions: {len(full_table)}, Resume point: {middle_idx}, Resume revid: {resume_revid}")
|
|
|
|
tester_partial = WikiqTester(SAILORMOON, "resume_partial", in_compression="7z", out_format="parquet")
|
|
partial_output_path = os.path.join(tester_partial.output, f"{SAILORMOON}.parquet")
|
|
|
|
partial_table = full_table.slice(0, middle_idx + 1)
|
|
pq.write_table(partial_table, partial_output_path)
|
|
|
|
try:
|
|
tester_partial.call_wikiq("--fandom-2020", "--resume")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
resumed_table = pq.read_table(partial_output_path)
|
|
|
|
resumed_df = resumed_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
full_df = full_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
|
|
assert_frame_equal(resumed_df, full_df, check_like=True, check_dtype=False)
|
|
|
|
print(f"Resume test passed! Original: {len(full_df)} rows, Resumed: {len(resumed_df)} rows")
|
|
|
|
|
|
def test_resume_with_diff():
|
|
"""Test that --resume works correctly with diff computation."""
|
|
tester_full = WikiqTester(SAILORMOON, "resume_diff_full", in_compression="7z", out_format="parquet")
|
|
|
|
try:
|
|
tester_full.call_wikiq("--diff", "--fandom-2020")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
full_output_path = os.path.join(tester_full.output, f"{SAILORMOON}.parquet")
|
|
full_table = pq.read_table(full_output_path)
|
|
|
|
resume_idx = len(full_table) // 3
|
|
resume_revid = full_table.column("revid")[resume_idx].as_py()
|
|
|
|
print(f"Total revisions: {len(full_table)}, Resume point: {resume_idx}, Resume revid: {resume_revid}")
|
|
|
|
tester_partial = WikiqTester(SAILORMOON, "resume_diff_partial", in_compression="7z", out_format="parquet")
|
|
partial_output_path = os.path.join(tester_partial.output, f"{SAILORMOON}.parquet")
|
|
|
|
partial_table = full_table.slice(0, resume_idx + 1)
|
|
pq.write_table(partial_table, partial_output_path)
|
|
|
|
try:
|
|
tester_partial.call_wikiq("--diff", "--fandom-2020", "--resume")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
resumed_table = pq.read_table(partial_output_path)
|
|
|
|
resumed_df = resumed_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
full_df = full_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
|
|
assert_frame_equal(resumed_df, full_df, check_like=True, check_dtype=False)
|
|
|
|
print(f"Resume with diff test passed! Original: {len(full_df)} rows, Resumed: {len(resumed_df)} rows")
|
|
|
|
|
|
def test_resume_with_partition_namespaces():
|
|
"""Test that --resume works correctly with --partition-namespaces.
|
|
|
|
Interrupts wikiq partway through processing, then resumes and verifies
|
|
the result matches an uninterrupted run. Uses --flush-per-batch to ensure
|
|
data is written to disk after each batch, making interruption deterministic.
|
|
"""
|
|
full_dir = os.path.join(TEST_OUTPUT_DIR, "resume_full")
|
|
partial_dir = os.path.join(TEST_OUTPUT_DIR, "resume_partial")
|
|
input_file = os.path.join(TEST_DIR, "dumps", f"{SAILORMOON}.xml.7z")
|
|
|
|
for output_dir in [full_dir, partial_dir]:
|
|
if os.path.exists(output_dir):
|
|
shutil.rmtree(output_dir)
|
|
os.makedirs(output_dir)
|
|
|
|
full_output = os.path.join(full_dir, f"{SAILORMOON}.parquet")
|
|
partial_output = os.path.join(partial_dir, f"{SAILORMOON}.parquet")
|
|
|
|
cmd_full = f"{WIKIQ} {input_file} -o {full_output} --batch-size 10 --partition-namespaces"
|
|
try:
|
|
subprocess.check_output(cmd_full, stderr=subprocess.PIPE, shell=True)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
full_dataset = ds.dataset(full_output, format="parquet", partitioning="hive")
|
|
full_df = full_dataset.to_table().to_pandas()
|
|
total_rows = len(full_df)
|
|
print(f"Full run produced {total_rows} rows")
|
|
|
|
batch_size = 10
|
|
cmd_partial = [
|
|
sys.executable, WIKIQ, input_file,
|
|
"-o", partial_output,
|
|
"--batch-size", str(batch_size),
|
|
"--partition-namespaces"
|
|
]
|
|
print(f"Starting: {' '.join(cmd_partial)}")
|
|
|
|
proc = subprocess.Popen(cmd_partial, stderr=subprocess.PIPE)
|
|
|
|
interrupt_delay = 5
|
|
time.sleep(interrupt_delay)
|
|
|
|
if proc.poll() is not None:
|
|
pytest.fail(f"wikiq completed in {interrupt_delay}s before we could interrupt")
|
|
|
|
print(f"Sending SIGUSR1 after {interrupt_delay}s")
|
|
proc.send_signal(signal.SIGUSR1)
|
|
|
|
try:
|
|
proc.wait(timeout=5)
|
|
print("Process exited gracefully after SIGUSR1")
|
|
except subprocess.TimeoutExpired:
|
|
print("Sending SIGTERM after SIGUSR1 timeout")
|
|
proc.send_signal(signal.SIGTERM)
|
|
proc.wait(timeout=30)
|
|
|
|
interrupted_dataset = ds.dataset(partial_output, format="parquet", partitioning="hive")
|
|
interrupted_rows = interrupted_dataset.count_rows()
|
|
print(f"Interrupted run wrote {interrupted_rows} rows")
|
|
|
|
assert interrupted_rows < total_rows, \
|
|
f"Process wrote all {interrupted_rows} rows before being killed"
|
|
|
|
cmd_resume = f"{WIKIQ} {input_file} -o {partial_output} --batch-size {batch_size} --partition-namespaces --resume"
|
|
try:
|
|
subprocess.check_output(cmd_resume, stderr=subprocess.PIPE, shell=True)
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
resumed_dataset = ds.dataset(partial_output, format="parquet", partitioning="hive")
|
|
resumed_df = resumed_dataset.to_table().to_pandas()
|
|
|
|
full_revids = set(full_df['revid'])
|
|
resumed_revids = set(resumed_df['revid'])
|
|
missing_revids = full_revids - resumed_revids
|
|
extra_revids = resumed_revids - full_revids
|
|
assert missing_revids == set() and extra_revids == set(), \
|
|
f"Revision ID mismatch: {len(missing_revids)} missing, {len(extra_revids)} extra. Missing: {sorted(missing_revids)[:10]}"
|
|
assert len(resumed_df) == len(full_df), \
|
|
f"Row count mismatch: {len(resumed_df)} vs {len(full_df)}"
|
|
|
|
print(f"Resume test passed! Full: {len(full_df)}, Interrupted: {interrupted_rows}, Resumed: {len(resumed_df)}")
|
|
|
|
|
|
def test_resume_file_not_found():
|
|
"""Test that --resume starts fresh when output file doesn't exist."""
|
|
tester = WikiqTester(SAILORMOON, "resume_not_found", in_compression="7z", out_format="parquet")
|
|
|
|
expected_output = os.path.join(tester.output, f"{SAILORMOON}.parquet")
|
|
if os.path.exists(expected_output):
|
|
os.remove(expected_output)
|
|
|
|
# Should succeed by starting fresh
|
|
tester.call_wikiq("--resume")
|
|
|
|
# Verify output was created
|
|
assert os.path.exists(expected_output), "Output file should be created when starting fresh"
|
|
table = pq.read_table(expected_output)
|
|
assert table.num_rows > 0, "Output should have data"
|
|
|
|
print("Resume file not found test passed - started fresh!")
|
|
|
|
|
|
def test_resume_simple():
|
|
"""Test that --resume works without --fandom-2020 and --partition-namespaces."""
|
|
tester_full = WikiqTester(SAILORMOON, "resume_simple_full", in_compression="7z", out_format="parquet")
|
|
|
|
try:
|
|
tester_full.call_wikiq()
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
full_output_path = os.path.join(tester_full.output, f"{SAILORMOON}.parquet")
|
|
full_table = pq.read_table(full_output_path)
|
|
|
|
resume_idx = len(full_table) // 3
|
|
resume_revid = full_table.column("revid")[resume_idx].as_py()
|
|
|
|
print(f"Total revisions: {len(full_table)}, Resume point: {resume_idx}, Resume revid: {resume_revid}")
|
|
|
|
tester_partial = WikiqTester(SAILORMOON, "resume_simple_partial", in_compression="7z", out_format="parquet")
|
|
partial_output_path = os.path.join(tester_partial.output, f"{SAILORMOON}.parquet")
|
|
|
|
partial_table = full_table.slice(0, resume_idx + 1)
|
|
pq.write_table(partial_table, partial_output_path)
|
|
|
|
try:
|
|
tester_partial.call_wikiq("--resume")
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
resumed_table = pq.read_table(partial_output_path)
|
|
|
|
resumed_df = resumed_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
full_df = full_table.to_pandas().sort_values("revid").reset_index(drop=True)
|
|
|
|
assert_frame_equal(resumed_df, full_df, check_like=True, check_dtype=False)
|
|
|
|
print(f"Resume simple test passed! Original: {len(full_df)} rows, Resumed: {len(resumed_df)} rows")
|
|
|
|
|
|
def test_resume_merge_with_invalid_temp_file():
|
|
"""Test that resume handles invalid/empty temp files gracefully.
|
|
|
|
This can happen when a namespace has no records after the resume point,
|
|
resulting in a temp file that was created but never written to.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
original_path = os.path.join(tmpdir, "original.parquet")
|
|
temp_path = os.path.join(tmpdir, "temp.parquet")
|
|
merged_path = os.path.join(tmpdir, "merged.parquet")
|
|
|
|
table = pa.table({"articleid": [1, 2, 3], "revid": [10, 20, 30]})
|
|
pq.write_table(table, original_path)
|
|
|
|
with open(temp_path, 'w') as f:
|
|
f.write("")
|
|
|
|
result = merge_parquet_files(original_path, temp_path, merged_path)
|
|
assert result == "original_only", f"Expected 'original_only' when temp file is invalid, got {result}"
|
|
|
|
assert os.path.exists(original_path), "Original file should still exist"
|
|
original_table = pq.read_table(original_path)
|
|
assert len(original_table) == 3, "Original file should be unchanged"
|
|
|
|
assert not os.path.exists(merged_path), "Merged file should not be created"
|
|
|
|
print("Resume merge with invalid temp file test passed!")
|
|
|
|
|
|
def test_resume_merge_with_corrupted_original():
|
|
"""Test that resume recovers from a corrupted original file if temp is valid.
|
|
|
|
This can happen if the original file was being written when the process
|
|
was killed, leaving it in a corrupted state.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
original_path = os.path.join(tmpdir, "original.parquet")
|
|
temp_path = os.path.join(tmpdir, "temp.parquet")
|
|
merged_path = os.path.join(tmpdir, "merged.parquet")
|
|
|
|
with open(original_path, 'w') as f:
|
|
f.write("corrupted data")
|
|
|
|
table = pa.table({"articleid": [4, 5, 6], "revid": [40, 50, 60]})
|
|
pq.write_table(table, temp_path)
|
|
|
|
result = merge_parquet_files(original_path, temp_path, merged_path)
|
|
assert result == "temp_only", f"Expected 'temp_only' when original is corrupted, got {result}"
|
|
|
|
assert not os.path.exists(merged_path), "Merged file should not be created for temp_only case"
|
|
|
|
print("Resume merge with corrupted original test passed!")
|
|
|
|
|
|
def test_resume_merge_both_invalid():
|
|
"""Test that resume handles both files being invalid."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
original_path = os.path.join(tmpdir, "original.parquet")
|
|
temp_path = os.path.join(tmpdir, "temp.parquet")
|
|
merged_path = os.path.join(tmpdir, "merged.parquet")
|
|
|
|
with open(original_path, 'w') as f:
|
|
f.write("corrupted original")
|
|
|
|
with open(temp_path, 'w') as f:
|
|
f.write("corrupted temp")
|
|
|
|
result = merge_parquet_files(original_path, temp_path, merged_path)
|
|
assert result == "both_invalid", f"Expected 'both_invalid' when both files corrupted, got {result}"
|
|
|
|
print("Resume merge with both invalid test passed!")
|
|
|
|
|
|
def test_cleanup_interrupted_resume_both_corrupted():
|
|
"""Test that cleanup_interrupted_resume returns 'start_fresh' when both files are corrupted."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
output_file = os.path.join(tmpdir, "output.parquet")
|
|
temp_file = output_file + ".resume_temp"
|
|
checkpoint_path = get_checkpoint_path(output_file, partition_namespaces=False)
|
|
|
|
with open(output_file, 'w') as f:
|
|
f.write("corrupted original")
|
|
|
|
with open(temp_file, 'w') as f:
|
|
f.write("corrupted temp")
|
|
|
|
with open(checkpoint_path, 'w') as f:
|
|
json.dump({"pageid": 100, "revid": 200}, f)
|
|
|
|
result = cleanup_interrupted_resume(output_file, partition_namespaces=False)
|
|
assert result == "start_fresh", f"Expected 'start_fresh', got {result}"
|
|
|
|
assert not os.path.exists(output_file), "Corrupted original should be deleted"
|
|
assert not os.path.exists(temp_file), "Corrupted temp should be deleted"
|
|
assert not os.path.exists(checkpoint_path), "Stale checkpoint should be deleted"
|
|
|
|
print("Cleanup interrupted resume with both corrupted test passed!")
|
|
|
|
|
|
def test_cleanup_interrupted_resume_original_corrupted_temp_valid():
|
|
"""Test that cleanup recovers from temp when original is corrupted."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
output_file = os.path.join(tmpdir, "output.parquet")
|
|
temp_file = output_file + ".resume_temp"
|
|
|
|
with open(output_file, 'w') as f:
|
|
f.write("corrupted original")
|
|
|
|
table = pa.table({"articleid": [10, 20, 30], "revid": [100, 200, 300]})
|
|
pq.write_table(table, temp_file)
|
|
|
|
result = cleanup_interrupted_resume(output_file, partition_namespaces=False)
|
|
assert result is None, f"Expected None (normal recovery), got {result}"
|
|
|
|
assert os.path.exists(output_file), "Output file should exist after recovery"
|
|
assert not os.path.exists(temp_file), "Temp file should be renamed to output"
|
|
|
|
recovered_table = pq.read_table(output_file)
|
|
assert len(recovered_table) == 3, "Recovered file should have 3 rows"
|
|
|
|
resume_point = get_resume_point(output_file, partition_namespaces=False)
|
|
assert resume_point is not None, "Should find resume point from recovered file"
|
|
assert resume_point == (30, 300, 0), f"Expected (30, 300, 0), got {resume_point}"
|
|
|
|
print("Cleanup with original corrupted, temp valid test passed!")
|
|
|
|
|
|
def test_cleanup_original_missing_temp_valid_no_checkpoint():
|
|
"""Test recovery when original is missing, temp is valid, and no checkpoint exists."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
output_file = os.path.join(tmpdir, "output.parquet")
|
|
temp_file = output_file + ".resume_temp"
|
|
checkpoint_path = get_checkpoint_path(output_file, partition_namespaces=False)
|
|
|
|
assert not os.path.exists(output_file)
|
|
|
|
table = pa.table({"articleid": [10, 20, 30], "revid": [100, 200, 300]})
|
|
pq.write_table(table, temp_file)
|
|
|
|
assert not os.path.exists(checkpoint_path)
|
|
|
|
result = cleanup_interrupted_resume(output_file, partition_namespaces=False)
|
|
assert result is None, f"Expected None (normal recovery), got {result}"
|
|
|
|
assert os.path.exists(output_file), "Output file should exist after recovery"
|
|
assert not os.path.exists(temp_file), "Temp file should be renamed to output"
|
|
|
|
resume_point = get_resume_point(output_file, partition_namespaces=False)
|
|
assert resume_point is not None, "Should find resume point from recovered file"
|
|
assert resume_point == (30, 300, 0), f"Expected (30, 300, 0), got {resume_point}"
|
|
|
|
print("Original missing, temp valid, no checkpoint test passed!")
|
|
|
|
|
|
def test_concurrent_jobs_different_input_files():
|
|
"""Test that merge only processes temp files for the current input file.
|
|
|
|
When multiple wikiq processes write to the same partitioned output directory
|
|
with different input files, each process should only merge its own temp files.
|
|
"""
|
|
from wikiq.resume import merge_partitioned_namespaces
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
# Create partitioned output structure
|
|
ns0_dir = os.path.join(tmpdir, "namespace=0")
|
|
ns1_dir = os.path.join(tmpdir, "namespace=1")
|
|
os.makedirs(ns0_dir)
|
|
os.makedirs(ns1_dir)
|
|
|
|
# Simulate two different input files producing output
|
|
file1 = "enwiki-20250123-pages-meta-history24-p1p100.parquet"
|
|
file2 = "enwiki-20250123-pages-meta-history24-p101p200.parquet"
|
|
|
|
# Create original and temp files for file1
|
|
table1_orig = pa.table({"articleid": [1, 2], "revid": [10, 20]})
|
|
table1_temp = pa.table({"articleid": [3, 4], "revid": [30, 40]})
|
|
pq.write_table(table1_orig, os.path.join(ns0_dir, file1))
|
|
pq.write_table(table1_temp, os.path.join(ns0_dir, file1 + ".resume_temp"))
|
|
pq.write_table(table1_orig, os.path.join(ns1_dir, file1))
|
|
pq.write_table(table1_temp, os.path.join(ns1_dir, file1 + ".resume_temp"))
|
|
|
|
# Create original and temp files for file2 (simulating another concurrent job)
|
|
table2_orig = pa.table({"articleid": [100, 200], "revid": [1000, 2000]})
|
|
table2_temp = pa.table({"articleid": [300, 400], "revid": [3000, 4000]})
|
|
pq.write_table(table2_orig, os.path.join(ns0_dir, file2))
|
|
pq.write_table(table2_temp, os.path.join(ns0_dir, file2 + ".resume_temp"))
|
|
pq.write_table(table2_orig, os.path.join(ns1_dir, file2))
|
|
pq.write_table(table2_temp, os.path.join(ns1_dir, file2 + ".resume_temp"))
|
|
|
|
# Merge only file1's temp files
|
|
merge_partitioned_namespaces(tmpdir, ".resume_temp", file1)
|
|
|
|
# Verify file1's temp files were merged and removed
|
|
assert not os.path.exists(os.path.join(ns0_dir, file1 + ".resume_temp")), \
|
|
"file1 temp should be merged in ns0"
|
|
assert not os.path.exists(os.path.join(ns1_dir, file1 + ".resume_temp")), \
|
|
"file1 temp should be merged in ns1"
|
|
|
|
# Verify file1's original now has merged data
|
|
merged1_ns0 = pq.read_table(os.path.join(ns0_dir, file1))
|
|
merged1_ns1 = pq.read_table(os.path.join(ns1_dir, file1))
|
|
assert merged1_ns0.num_rows == 4, f"file1 ns0 should have 4 rows after merge, got {merged1_ns0.num_rows}"
|
|
assert merged1_ns1.num_rows == 4, f"file1 ns1 should have 4 rows after merge, got {merged1_ns1.num_rows}"
|
|
|
|
# Verify file2's temp files are UNTOUCHED (still exist)
|
|
assert os.path.exists(os.path.join(ns0_dir, file2 + ".resume_temp")), \
|
|
"file2 temp should NOT be touched in ns0"
|
|
assert os.path.exists(os.path.join(ns1_dir, file2 + ".resume_temp")), \
|
|
"file2 temp should NOT be touched in ns1"
|
|
|
|
# Verify file2's original is unchanged
|
|
orig2_ns0 = pq.read_table(os.path.join(ns0_dir, file2))
|
|
orig2_ns1 = pq.read_table(os.path.join(ns1_dir, file2))
|
|
assert orig2_ns0.num_rows == 2, "file2 ns0 should still have 2 rows"
|
|
assert orig2_ns1.num_rows == 2, "file2 ns1 should still have 2 rows"
|
|
|
|
print("Concurrent jobs with different input files test passed!")
|
|
|
|
|
|
def test_max_revisions_per_file_creates_parts():
|
|
"""Test that --max-revisions-per-file creates multiple part files."""
|
|
import re
|
|
tester = WikiqTester(SAILORMOON, "max_revs_parts", in_compression="7z", out_format="parquet")
|
|
|
|
max_revs = 50
|
|
try:
|
|
# Use a very small limit to force multiple parts
|
|
tester.call_wikiq("--fandom-2020", "--max-revisions-per-file", str(max_revs))
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
output_dir = tester.output
|
|
all_parquet = [f for f in os.listdir(output_dir) if f.endswith(".parquet") and ".part" in f]
|
|
|
|
# Sort by part number numerically
|
|
def get_part_num(filename):
|
|
match = re.search(r'\.part(\d+)\.parquet$', filename)
|
|
return int(match.group(1)) if match else 0
|
|
|
|
part_files = sorted(all_parquet, key=get_part_num)
|
|
|
|
assert len(part_files) > 1, f"Expected multiple part files, got {part_files}"
|
|
|
|
# Read all parts and verify total rows
|
|
total_rows = 0
|
|
for f in part_files:
|
|
table = pq.read_table(os.path.join(output_dir, f))
|
|
total_rows += len(table)
|
|
|
|
assert total_rows > 0, "Should have some rows across all parts"
|
|
|
|
# Each part (except the last) should have at least max_revisions rows
|
|
# (rotation happens after the batch that hits the limit is written)
|
|
for f in part_files[:-1]:
|
|
table = pq.read_table(os.path.join(output_dir, f))
|
|
assert len(table) >= max_revs, f"Part file {f} should have at least {max_revs} rows, got {len(table)}"
|
|
|
|
print(f"max-revisions-per-file test passed! Created {len(part_files)} parts with {total_rows} total rows")
|
|
|
|
|
|
def test_max_revisions_per_file_with_partitioned():
|
|
"""Test that --max-revisions-per-file works with partitioned namespace output."""
|
|
import re
|
|
tester = WikiqTester(SAILORMOON, "max_revs_partitioned", in_compression="7z", out_format="parquet")
|
|
|
|
max_revs = 20
|
|
try:
|
|
# Use a small limit to force parts, with partitioned output
|
|
tester.call_wikiq("--fandom-2020", "--partition-namespaces", "--max-revisions-per-file", str(max_revs))
|
|
except subprocess.CalledProcessError as exc:
|
|
pytest.fail(exc.stderr.decode("utf8"))
|
|
|
|
output_dir = tester.output
|
|
|
|
# Find namespace directories
|
|
ns_dirs = [d for d in os.listdir(output_dir) if d.startswith("namespace=")]
|
|
assert len(ns_dirs) > 0, "Should have namespace directories"
|
|
|
|
def get_part_num(filename):
|
|
match = re.search(r'\.part(\d+)\.parquet$', filename)
|
|
return int(match.group(1)) if match else 0
|
|
|
|
# Check that at least one namespace has multiple parts
|
|
found_multi_part = False
|
|
for ns_dir in ns_dirs:
|
|
ns_path = os.path.join(output_dir, ns_dir)
|
|
parquet_files = [f for f in os.listdir(ns_path) if f.endswith(".parquet")]
|
|
part_files = [f for f in parquet_files if ".part" in f]
|
|
if len(part_files) > 1:
|
|
found_multi_part = True
|
|
# Sort by part number and verify each part (except last) has at least limit rows
|
|
sorted_parts = sorted(part_files, key=get_part_num)
|
|
for f in sorted_parts[:-1]:
|
|
pf = pq.ParquetFile(os.path.join(ns_path, f))
|
|
num_rows = pf.metadata.num_rows
|
|
assert num_rows >= max_revs, f"Part file {f} in {ns_dir} should have at least {max_revs} rows, got {num_rows}"
|
|
|
|
assert found_multi_part, "At least one namespace should have multiple part files"
|
|
|
|
print(f"max-revisions-per-file with partitioned output test passed!")
|