diff --git a/compass/pipeline/collection/persistence.py b/compass/pipeline/collection/persistence.py index 58c844738..e3d54209b 100644 --- a/compass/pipeline/collection/persistence.py +++ b/compass/pipeline/collection/persistence.py @@ -1,8 +1,11 @@ """Persistence for collected documents""" +import os import json import asyncio from pathlib import Path +from glob import glob +from itertools import chain from statistics import median from collections import Counter from warnings import warn @@ -123,13 +126,15 @@ async def write_collection_manifest_shard(shard_dir, collection_info): ) -async def load_collection_manifest(manifest_fp, expected_tech): - """Load a collection manifest from disk +async def load_collection_manifest_jurisdictions(manifest_fp, expected_tech): + """Load jurisdictions from one or more collection manifest(s) Parameters ---------- - manifest_fp : path-like - Path to the collection manifest file to be loaded. + manifest_fp : path-like or list of path-like + Path to the collection manifest file to be loaded. Can be a + single path or a list of paths, any of which may include glob + patterns. expected_tech : str Technology specified in the pipeline request, used to validate compatibility with the manifest. @@ -137,11 +142,45 @@ async def load_collection_manifest(manifest_fp, expected_tech): Returns ------- dict - Loaded collection manifest as a dictionary. + Mapping of FIPS codes to jurisdiction infos from the collection + manifest(s). + + Raises + ------ + COMPASSValueError + If a duplicate jurisdiction is found in the manifest(s). """ - return await GenericFuncRunner.call( - _load_collection_manifest, manifest_fp, expected_tech - ) + if isinstance(manifest_fp, (str, os.PathLike)): + manifest_fp = [str(manifest_fp)] + + task_fps = [] + for maybe_glob in manifest_fp: + # ruff: ignore[glob] + new_fps = [ + Path(match) for match in glob(str(maybe_glob), recursive=True) + ] + task_fps.extend(new_fps or [maybe_glob]) + + tasks = [ + GenericFuncRunner.call(_load_collection_manifest, fp, expected_tech) + for fp in task_fps + ] + manifests = await asyncio.gather(*tasks) + + jurisdictions_by_fips = {} + for jurisdiction in chain.from_iterable( + manifest.get("jurisdictions", []) for manifest in manifests + ): + if jurisdiction is None: + continue + + fips = jurisdiction.get("FIPS") + if fips in jurisdictions_by_fips: + msg = f"Duplicate collection manifest entry for FIPS '{fips}'" + raise COMPASSValueError(msg) + jurisdictions_by_fips[fips] = jurisdiction + + return jurisdictions_by_fips async def load_specific_collection_manifest_shard(shard_dir, jurisdiction): @@ -191,7 +230,9 @@ def _write_collection_manifest_shard(shard_dir, collection_info): def _load_collection_manifest(manifest_fp, expected_tech): """Load a collection manifest from disk""" try: - manifest = load_config(manifest_fp, file_name="Collection manifest") + manifest = load_config( + manifest_fp, resolve_paths=True, file_name="Collection manifest" + ) except COMPASSFileNotFoundError: manifest = _load_collection_manifest_from_shards( manifest_fp, expected_tech @@ -223,6 +264,8 @@ def _load_specific_collection_manifest_shard(shard_dir, jurisdiction): return load_config( shard_fp, + # paths are NOT relative to the shard directory, so should not + # be resolved here resolve_paths=False, file_name="Collection manifest shard", ) @@ -421,6 +464,9 @@ def _load_collection_manifest_from_shards(manifest_fp, expected_tech): for shard_fp in shard_fps: collection_info = load_config( shard_fp, + # paths are NOT relative to the shard directory, so should + # not be resolved here; they are resolved using the + # `resolve_all_paths` function call below resolve_paths=False, file_name="Collection manifest shard", ) diff --git a/compass/pipeline/coordinator.py b/compass/pipeline/coordinator.py index 7b890d3d7..16956fab5 100644 --- a/compass/pipeline/coordinator.py +++ b/compass/pipeline/coordinator.py @@ -25,7 +25,7 @@ from compass.pipeline.collection.persistence import ( build_collection_manifest, write_collection_manifest, - load_collection_manifest, + load_collection_manifest_jurisdictions, ) from compass.pipeline import BaseRequest from compass.pipeline.runtime import PipelineRuntime @@ -300,10 +300,12 @@ async def run(self, jurisdictions_df): logger.debug( "Manifest path(s): %s", self.runtime.request.collection_manifest_fp ) - manifest = await load_collection_manifest( - self.runtime.request.collection_manifest_fp, self.runtime.tech + collection_infos_by_fips = ( + await load_collection_manifest_jurisdictions( + self.runtime.request.collection_manifest_fp, self.runtime.tech + ) ) - jurisdictions = manifest.get("jurisdictions", []) + logger.info( "Extracting structured data for %d jurisdiction(s)", len(jurisdictions_df), @@ -312,12 +314,8 @@ async def run(self, jurisdictions_df): tasks = [] start_date = datetime.now(UTC) for jurisdiction in jurisdictions_from_df(jurisdictions_df): - collection_info = [ - info - for info in jurisdictions - if info is not None and info.get("FIPS") == jurisdiction.code - ] - if not collection_info: + collection_info = collection_infos_by_fips.get(jurisdiction.code) + if collection_info is None: logger.warning( "No collection info found for %s; skipping extraction", jurisdiction.full_name, @@ -330,7 +328,7 @@ async def run(self, jurisdictions_df): workflow = self._create(jurisdiction, usage_tracker=usage_tracker) tasks.append( asyncio.create_task( - workflow.run_extraction_with_logging(collection_info[0]), + workflow.run_extraction_with_logging(collection_info), name=jurisdiction.full_name, ) ) diff --git a/compass/pipeline/data_classes.py b/compass/pipeline/data_classes.py index 9f3afbf47..278703ce5 100644 --- a/compass/pipeline/data_classes.py +++ b/compass/pipeline/data_classes.py @@ -624,14 +624,16 @@ def __init__( # ruff:ignore[too-many-arguments] terminal. If ``True``, all of the unordered records are written to a "all.log" file in the `log_dir` directory. By default, ``False``. - collection_manifest_fp : path-like, optional + collection_manifest_fp : path-like or list of path-like, optional Path to the JSON collection manifest created by the document - collection step. The manifest must contain the persisted - document information needed to reload each collected - document for extraction. Only needed if running in + collection step. This can be a single path or a list of + paths for multiple collection manifests, any of which may + include glob patterns. Each collection manifest must contain + the persisted document information needed to reload each + collected document for extraction. Only needed if running in extraction mode with a separate collection step. By default, ``None``. - """ + """ # ruff:ignore[doc-line-too-long] self.tech = tech self.jurisdiction_fp = jurisdiction_fp self.perform_se_search = perform_se_search @@ -1092,11 +1094,13 @@ def __init__( # ruff:ignore[too-many-arguments] name of the subdivision, and the "Jurisdiction Type" should be a string identifying the type of subdivision (e.g., "City", "Township", etc.) - collection_manifest_fp : path-like + collection_manifest_fp : path-like or list of path-like, optional Path to the JSON collection manifest created by the document - collection step. The manifest must contain the persisted - document information needed to reload each collected - document for extraction. + collection step. This can be a single path or a list of + paths for multiple collection manifests, any of which may + include glob patterns. Each collection manifest must contain + the persisted document information needed to reload each + collected document for extraction. By default, ``None``. model : str or list of dict, default="gpt-4o-mini" LLM model(s) to use for scraping and parsing ordinance documents. If a string is provided, it is assumed to be the @@ -1228,7 +1232,7 @@ def __init__( # ruff:ignore[too-many-arguments] terminal. If ``True``, all of the unordered records are written to a "all.log" file in the `log_dir` directory. By default, ``False``. - """ + """ # ruff:ignore[doc-line-too-long] super().__init__( out_dir=out_dir, diff --git a/tests/python/unit/pipeline/test_pipeline_collection_persistence.py b/tests/python/unit/pipeline/test_pipeline_collection_persistence.py index 99321aed6..853c531e8 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection_persistence.py +++ b/tests/python/unit/pipeline/test_pipeline_collection_persistence.py @@ -8,7 +8,10 @@ import pytest import compass.pipeline.collection.persistence as persistence_module +from compass.exceptions import COMPASSValueError from compass.pipeline.collection.dedupe import DocumentDeDuplicator +from compass.services.provider import RunningAsyncServices +from compass.services.threaded import GenericFuncRunner def _build_doc(source, pages, *, has_parsed_text=True): @@ -29,6 +32,231 @@ def _build_jurisdiction(full_name="Example Township", code="12345"): return SimpleNamespace(full_name=full_name, code=code) +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["str", "path"]) +@pytest.mark.parametrize("is_relative", [True, False]) +@pytest.mark.parametrize("has_wildcard", [True, False]) +@pytest.mark.parametrize("is_list", [True, False]) +# ruff:ignore[complex-structure] +async def test_load_collection_manifest_jurisdictions_path_variants( + tmp_path, monkeypatch, input_type, is_relative, has_wildcard, is_list +): + """Manifest inputs and persisted document paths should resolve""" + manifest_dir = tmp_path / "manifests" + manifest_fps = [ + manifest_dir / "first" / "manifest_first.json", + manifest_dir / "second" / "manifest_second.json", + ] + expected_jurisdictions = [] + for index, manifest_fp in enumerate(manifest_fps, start=1): + document_paths = { + "dot": "./documents/source.html", + "parent": "../shared/source.html", + "normalized": "./documents/../normalized/source.html", + "windows_dot": r".\documents\source.html", + "windows_parent": r"..\shared\source.html", + } + documents = [ + { + "path_case": path_case, + "source_fp": source_fp, + "parsed_fp": source_fp.replace("source.html", "parsed.txt"), + } + for path_case, source_fp in document_paths.items() + ] + jurisdiction = {"FIPS": f"{index:03d}", "documents": documents} + manifest_fp.parent.mkdir(parents=True) + manifest_fp.write_text( + json.dumps( + { + "tech": "solar", + "jurisdictions": [jurisdiction], + } + ), + encoding="utf-8", + ) + expected_jurisdictions.append( + { + "FIPS": f"{index:03d}", + "documents": [ + { + "path_case": doc_info["path_case"], + "source_fp": str( + ( + manifest_fp.parent + / doc_info["source_fp"].replace("\\", "/") + ) + .resolve() + .as_posix() + ), + "parsed_fp": str( + ( + manifest_fp.parent + / doc_info["parsed_fp"].replace("\\", "/") + ) + .resolve() + .as_posix() + ), + } + for doc_info in documents + ], + } + ) + + manifest_inputs = [] + for manifest_fp in manifest_fps: + if is_relative: + manifest_input = f"./{manifest_fp.relative_to(tmp_path)}" + else: + manifest_input = manifest_fp + if has_wildcard: + manifest_input = str(manifest_input).replace( + manifest_fp.name, "*.json" + ) + if input_type == "path": + manifest_input = Path(manifest_input) + else: + manifest_input = str(manifest_input) + manifest_inputs.append(manifest_input) + + monkeypatch.chdir(tmp_path) + manifest_input = manifest_inputs if is_list else manifest_inputs[0] + async with RunningAsyncServices([GenericFuncRunner()]): + jurisdictions = ( + await persistence_module.load_collection_manifest_jurisdictions( + manifest_input, "solar" + ) + ) + + if not is_list: + expected_jurisdictions = expected_jurisdictions[:1] + expected_jurisdictions = { + jurisdiction["FIPS"]: jurisdiction + for jurisdiction in expected_jurisdictions + } + assert jurisdictions == expected_jurisdictions + for fips, jurisdiction in sorted(jurisdictions.items()): + manifest_fp = manifest_fps[int(fips) - 1] + for doc_info in jurisdiction["documents"]: + for key in ("source_fp", "parsed_fp"): + assert Path(doc_info[key]).is_absolute() + expected_path = document_paths[doc_info["path_case"]] + if key == "parsed_fp": + expected_path = expected_path.replace( + "source.html", "parsed.txt" + ) + expected_path = expected_path.replace("\\", "/") + assert doc_info[key] == str( + (manifest_fp.parent / expected_path).resolve().as_posix() + ) + + +@pytest.mark.asyncio +async def test_load_collection_manifest_jurisdictions_recursive_wildcard( + tmp_path, +): + """Recursive wildcard inputs should load nested manifests""" + manifest_dir = tmp_path / "manifests" + manifest_fps = [ + manifest_dir / "first" / "collection_manifest.json", + manifest_dir / "second" / "nested" / "collection_manifest.json", + ] + for index, manifest_fp in enumerate(manifest_fps, start=1): + manifest_fp.parent.mkdir(parents=True) + manifest_fp.write_text( + json.dumps( + { + "tech": "solar", + "jurisdictions": [{"FIPS": f"{index:03d}"}], + } + ), + encoding="utf-8", + ) + + async with RunningAsyncServices([GenericFuncRunner()]): + jurisdictions = ( + await persistence_module.load_collection_manifest_jurisdictions( + manifest_dir / "**" / "*.json", "solar" + ) + ) + + assert jurisdictions == {"001": {"FIPS": "001"}, "002": {"FIPS": "002"}} + + +@pytest.mark.asyncio +async def test_load_collection_manifest_jurisdictions_rejects_duplicate_fips( + tmp_path, +): + """Overlapping manifests should fail instead of discarding an entry""" + manifest_fps = [tmp_path / "first.json", tmp_path / "second.json"] + for manifest_fp in manifest_fps: + manifest_fp.write_text( + json.dumps( + { + "tech": "solar", + "jurisdictions": [{"FIPS": "12345", "documents": []}], + } + ), + encoding="utf-8", + ) + + async with RunningAsyncServices([GenericFuncRunner()]): + with pytest.raises( + COMPASSValueError, + match="Duplicate collection manifest entry for FIPS '12345'", + ): + await persistence_module.load_collection_manifest_jurisdictions( + manifest_fps, "solar" + ) + + +@pytest.mark.asyncio +async def test_load_collection_manifest_jurisdictions_resolves_shard_paths( + tmp_path, +): + """Shard-recovered document paths should resolve from manifest root""" + manifest_dir = tmp_path / "collection" + shard_dir = manifest_dir / "shards" + shard_dir.mkdir(parents=True) + collection_info = { + "FIPS": "12345", + "full_name": "Example Township", + "documents": [ + { + "source_fp": "./downloaded/source.html", + "parsed_fp": "./parsed/source.txt", + } + ], + } + shard_fp = ( + shard_dir + / persistence_module._collection_manifest_shard_filename( + collection_info + ) + ) + shard_fp.write_text(json.dumps(collection_info), encoding="utf-8") + + manifest_fp = ( + manifest_dir / persistence_module.COLLECTION_MANIFEST_FILENAME + ) + async with RunningAsyncServices([GenericFuncRunner()]): + jurisdictions = ( + await persistence_module.load_collection_manifest_jurisdictions( + manifest_fp, "solar" + ) + ) + + document = jurisdictions["12345"]["documents"][0] + assert document["source_fp"] == str( + (manifest_dir / "downloaded/source.html").resolve().as_posix() + ) + assert document["parsed_fp"] == str( + (manifest_dir / "parsed/source.txt").resolve().as_posix() + ) + assert Path(document["source_fp"]).is_absolute() + assert Path(document["parsed_fp"]).is_absolute() + + @pytest.mark.asyncio async def test_persist_documents_filters_docs_without_parsed_text( monkeypatch, tmp_path