mirror of
https://github.com/migatu/conjurer.git
synced 2026-07-14 21:38:38 +00:00
Fixes
This commit is contained in:
@@ -0,0 +1,193 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
import argparse
|
||||||
|
import logging
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
# =========================
|
||||||
|
# Configuration constants
|
||||||
|
# =========================
|
||||||
|
|
||||||
|
MAX_CHUNK_SIZE_BYTES = 50 * 1024 * 1024 # 100 MB
|
||||||
|
OUTPUT_SUFFIX = "_chunk.txt"
|
||||||
|
|
||||||
|
|
||||||
|
# =========================
|
||||||
|
# Logging
|
||||||
|
# =========================
|
||||||
|
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO,
|
||||||
|
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def get_default_output_dir(input_file: Path) -> Path:
|
||||||
|
"""
|
||||||
|
Return default output directory placed next to the input file.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
/data/big.txt -> /data/big/
|
||||||
|
"""
|
||||||
|
return input_file.with_suffix("")
|
||||||
|
|
||||||
|
|
||||||
|
def write_chunk(output_dir: Path, chunk_index: int, lines: list[bytes], chunk_size: int) -> Path:
|
||||||
|
"""
|
||||||
|
Write one chunk file and return its path.
|
||||||
|
"""
|
||||||
|
output_file = output_dir / f"{chunk_index}{OUTPUT_SUFFIX}"
|
||||||
|
|
||||||
|
with output_file.open("wb") as file:
|
||||||
|
file.writelines(lines)
|
||||||
|
|
||||||
|
actual_size = output_file.stat().st_size
|
||||||
|
|
||||||
|
if actual_size != chunk_size:
|
||||||
|
logger.warning(
|
||||||
|
"Chunk size mismatch for %s: expected %d bytes, got %d bytes",
|
||||||
|
output_file,
|
||||||
|
chunk_size,
|
||||||
|
actual_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
return output_file
|
||||||
|
|
||||||
|
|
||||||
|
def split_txt_file(input_file: Path, output_dir: Path | None = None) -> None:
|
||||||
|
"""
|
||||||
|
Split a text-like file into smaller files limited by MAX_CHUNK_SIZE_BYTES.
|
||||||
|
|
||||||
|
The file is processed in binary mode, so invalid UTF-8 or mixed encodings
|
||||||
|
do not break the split. Lines are never split.
|
||||||
|
"""
|
||||||
|
if not input_file.exists():
|
||||||
|
raise FileNotFoundError(f"Input file does not exist: {input_file}")
|
||||||
|
|
||||||
|
if not input_file.is_file():
|
||||||
|
raise ValueError(f"Input path is not a file: {input_file}")
|
||||||
|
|
||||||
|
if MAX_CHUNK_SIZE_BYTES <= 0:
|
||||||
|
raise ValueError("MAX_CHUNK_SIZE_BYTES must be greater than zero")
|
||||||
|
|
||||||
|
if output_dir is None:
|
||||||
|
output_dir = get_default_output_dir(input_file)
|
||||||
|
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
logger.info("Input file: %s", input_file)
|
||||||
|
logger.info("Output directory: %s", output_dir)
|
||||||
|
logger.info("Max chunk size: %d bytes", MAX_CHUNK_SIZE_BYTES)
|
||||||
|
|
||||||
|
chunk_index = 0
|
||||||
|
current_lines: list[bytes] = []
|
||||||
|
current_size = 0
|
||||||
|
total_lines = 0
|
||||||
|
total_bytes = 0
|
||||||
|
|
||||||
|
with input_file.open("rb") as input_handle:
|
||||||
|
for line_number, line in enumerate(input_handle, start=1):
|
||||||
|
line_size = len(line)
|
||||||
|
total_lines = line_number
|
||||||
|
total_bytes += line_size
|
||||||
|
|
||||||
|
if line_size > MAX_CHUNK_SIZE_BYTES:
|
||||||
|
logger.warning(
|
||||||
|
"Line %d is larger than max chunk size: %d bytes > %d bytes. "
|
||||||
|
"It will be written as a separate chunk.",
|
||||||
|
line_number,
|
||||||
|
line_size,
|
||||||
|
MAX_CHUNK_SIZE_BYTES,
|
||||||
|
)
|
||||||
|
|
||||||
|
if current_lines and current_size + line_size > MAX_CHUNK_SIZE_BYTES:
|
||||||
|
output_file = write_chunk(
|
||||||
|
output_dir=output_dir,
|
||||||
|
chunk_index=chunk_index,
|
||||||
|
lines=current_lines,
|
||||||
|
chunk_size=current_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Written chunk %d: %s (%d bytes)",
|
||||||
|
chunk_index,
|
||||||
|
output_file,
|
||||||
|
current_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
chunk_index += 1
|
||||||
|
current_lines = []
|
||||||
|
current_size = 0
|
||||||
|
|
||||||
|
current_lines.append(line)
|
||||||
|
current_size += line_size
|
||||||
|
|
||||||
|
if current_lines:
|
||||||
|
output_file = write_chunk(
|
||||||
|
output_dir=output_dir,
|
||||||
|
chunk_index=chunk_index,
|
||||||
|
lines=current_lines,
|
||||||
|
chunk_size=current_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"Written chunk %d: %s (%d bytes)",
|
||||||
|
chunk_index,
|
||||||
|
output_file,
|
||||||
|
current_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
chunk_count = chunk_index + 1
|
||||||
|
else:
|
||||||
|
logger.info("Input file is empty. No chunks were created.")
|
||||||
|
chunk_count = 0
|
||||||
|
|
||||||
|
logger.info("Total lines processed: %d", total_lines)
|
||||||
|
logger.info("Total bytes processed: %d", total_bytes)
|
||||||
|
logger.info("Chunks created: %d", chunk_count)
|
||||||
|
logger.info("Done.")
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Split a large TXT-like file into smaller line-safe chunks."
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"input_file",
|
||||||
|
type=Path,
|
||||||
|
help="Path to input TXT file.",
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"-o",
|
||||||
|
"--output-dir",
|
||||||
|
type=Path,
|
||||||
|
default=None,
|
||||||
|
help=(
|
||||||
|
"Optional output directory. "
|
||||||
|
"By default, a directory named after the input file is created next to it."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
args = parse_args()
|
||||||
|
|
||||||
|
try:
|
||||||
|
split_txt_file(args.input_file, args.output_dir)
|
||||||
|
return 0
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Failed to split file")
|
||||||
|
return 1
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
sys.exit(main())
|
||||||
|
|
||||||
@@ -19,7 +19,7 @@ from requests import ConnectTimeout, Timeout
|
|||||||
import lib_paths
|
import lib_paths
|
||||||
|
|
||||||
SCR_DATABASE_PATH = os.getenv("CONJURER_LIBRARIAN_DB_PATH", r"C:\\Database\\chunks\\")
|
SCR_DATABASE_PATH = os.getenv("CONJURER_LIBRARIAN_DB_PATH", r"C:\\Database\\chunks\\")
|
||||||
SCR_FILENAME = os.getenv("CONJURER_LIBRARIAN_SCRAPE_CHUNK", "40_chunk.txt")
|
SCR_FILENAME = os.getenv("CONJURER_LIBRARIAN_SCRAPE_CHUNK", "44_chunk.txt")
|
||||||
SCR_ENCODING = os.getenv("CONJURER_ENCODING", "utf-8")
|
SCR_ENCODING = os.getenv("CONJURER_ENCODING", "utf-8")
|
||||||
|
|
||||||
WORK_Q = Queue()
|
WORK_Q = Queue()
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ services:
|
|||||||
volumes:
|
volumes:
|
||||||
# Local DOI chunk database (0_chunk.txt ... N_chunk.txt). Read-only:
|
# Local DOI chunk database (0_chunk.txt ... N_chunk.txt). Read-only:
|
||||||
# the search workers only read it (open mode "r").
|
# the search workers only read it (open mode "r").
|
||||||
- /srv/librarian/doi:/doi:ro
|
- /mnt/conjurer_swap/librarian_data/base_it1:/doi:ro
|
||||||
# Runtime JSON state (cr_results/rr_results/not_in_db/s_results) -
|
# Runtime JSON state (cr_results/rr_results/not_in_db/s_results) -
|
||||||
# seeded on first run, persisted here across restarts.
|
# seeded on first run, persisted here across restarts.
|
||||||
- /srv/librarian/state:/lib_temp_files
|
- /srv/librarian/state:/lib_temp_files
|
||||||
|
|||||||
Reference in New Issue
Block a user