diff --git a/CHANGELOG.md b/CHANGELOG.md index 2700790..ae82797 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,15 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] +### Added + +- `PerFileConverterMixin`: converts each source of a converter on its own and merges the parts, + for sources that do not fit into memory together. + ### Fixed - Longer HTTP download limit, 5 minute can be too short +- `--compression none` no longer fails in the DuckDB converter and in `vec merge` with DuckDB. ## [v0.4.1] - 2026-10-25 diff --git a/tests/test_convert_duckdb.py b/tests/test_convert_duckdb.py index ac6fa44..2d12c27 100644 --- a/tests/test_convert_duckdb.py +++ b/tests/test_convert_duckdb.py @@ -10,6 +10,7 @@ import shapely from loguru import logger +from vecorel_cli.const import COMPRESSION_METHODS from vecorel_cli.conversion import duckdb as duckdb_module from vecorel_cli.conversion.base import BaseConverter from vecorel_cli.conversion.duckdb import DuckDBBaseConverter, _constants_table @@ -563,6 +564,20 @@ def test_merge_parquet(tmp_folder): assert ValidateData().validate(dest, num=100, schema_map={}).errors == [] +@pytest.mark.parametrize("compression", COMPRESSION_METHODS) +def test_every_compression_method_is_written(tmp_folder, compression): + src = _source_file(tmp_folder) + part = tmp_folder / "part.parquet" + Converter().convert(part, input_files={src: "source.parquet"}, compression=compression) + dest = tmp_folder / "merged.parquet" + Converter().merge_parquet([part], dest, compression=compression) + + for path in (part, dest): + with pq.ParquetFile(path) as pf: + codec = pf.metadata.row_group(0).column(0).compression + assert codec == ("UNCOMPRESSED" if compression == "none" else compression.upper()) + + def test_merge_parquet_sees_an_id_that_repeats_across_parts(tmp_folder, capsys): """A per-part check cannot see this; a check over the merge can.""" parts = [] diff --git a/tests/test_per_file.py b/tests/test_per_file.py new file mode 100644 index 0000000..802c0b3 --- /dev/null +++ b/tests/test_per_file.py @@ -0,0 +1,95 @@ +import io +import sys + +import geopandas as gpd +import multivolumefile +import py7zr +import shapely +from loguru import logger + +from vecorel_cli.conversion.base import BaseConverter +from vecorel_cli.conversion.per_file import PerFileConverterMixin + + +class Converter(PerFileConverterMixin, BaseConverter): + id = "test" + short_name = "Test" + title = "Test dataset" + description = "Test dataset" + license = "CC0-1.0" + variants = {"2024": None, "2025": None} + columns = {"geometry": "geometry", "id": "id", "region": "region"} + missing_schemas = {"properties": {"region": {"type": "string"}}} + + +def _source(folder, region): + gdf = gpd.GeoDataFrame( + { + "id": [f"{region}-1", f"{region}-2"], + "region": [region, region], + "geometry": [shapely.box(0, 0, 1, 1), shapely.box(2, 0, 3, 1)], + }, + crs="EPSG:4326", + ) + path = folder / f"{region}.parquet" + gdf.to_parquet(path) + return path + + +def _capture_log(): + # loguru's default sink is bound to stderr at import, so point it at stdout + logger.remove() + logger.add(sys.stdout, format="{message}", level="DEBUG", colorize=False) + + +def test_each_source_is_converted_on_its_own_and_merged(tmp_path, capsys): + sources = [_source(tmp_path, region) for region in ("a", "b")] + output = tmp_path / "out" / "merged.parquet" + _capture_log() + + Converter().convert(str(output), input_files={str(s): s.name for s in sources}) + + merged = gpd.read_parquet(output) + assert sorted(merged["id"]) == ["a-1", "a-2", "b-1", "b-2"] + # constant within each part, so it only survives the merge as a column + assert sorted(merged["region"].unique()) == ["a", "b"] + assert [p.name for p in output.parent.iterdir()] == ["merged.parquet"] # the parts are gone + + log = capsys.readouterr().out + assert log.count("Using user provided input file(s)") == 1 + assert log.count("No variant given") == 1 + + +def test_a_single_source_is_converted_directly(tmp_path): + source = _source(tmp_path, "a") + output = tmp_path / "single.parquet" + + Converter().convert(str(output), input_files={str(source): source.name}) + + assert sorted(gpd.read_parquet(output)["id"]) == ["a-1", "a-2"] + + +def test_the_volumes_of_a_7z_archive_are_converted_together(tmp_path): + archived = _source(tmp_path, "b") + measure = io.BytesIO() + with py7zr.SevenZipFile(measure, "w") as archive: + archive.write(archived, archived.name) + volume_size = len(measure.getvalue()) // 3 + 1 + with multivolumefile.MultiVolume( + str(tmp_path / "b.7z"), mode="wb", volume=volume_size, ext_digits=3 + ) as volume: + with py7zr.SevenZipFile(volume, "w") as archive: + archive.write(archived, archived.name) + volumes = sorted(tmp_path.glob("b.7z.*")) + assert len(volumes) > 1, "the archive must span several volumes for this test" + archived.unlink() + + source = _source(tmp_path, "a") + output = tmp_path / "merged.parquet" + # only the first volume names the target file + input_files = {str(source): source.name} + input_files |= {str(v): ([archived.name] if i == 0 else []) for i, v in enumerate(volumes)} + + Converter().convert(str(output), input_files=input_files) + + assert sorted(gpd.read_parquet(output)["id"]) == ["a-1", "a-2", "b-1", "b-2"] diff --git a/vecorel_cli/conversion/base.py b/vecorel_cli/conversion/base.py index 6fca35f..8b5fbf6 100644 --- a/vecorel_cli/conversion/base.py +++ b/vecorel_cli/conversion/base.py @@ -516,6 +516,15 @@ def get_urls(self): raise ValueError(f"Unknown variant '{self.variant}', choose from {opts}") return urls + def get_input_urls(self, input_files=None): + if input_files is not None and isinstance(input_files, dict) and len(input_files) > 0: + self.warning("Using user provided input file(s) instead of the pre-defined file(s)") + return input_files + urls = self.get_urls() + if urls is None: + raise ValueError("No input files provided") + return urls + def get_data( self, paths: list[tuple[str, str]], **kwargs ) -> Generator[tuple[GeoDataFrame, str, str, Optional[str]]]: @@ -658,13 +667,7 @@ def convert( if directory: os.makedirs(directory, exist_ok=True) - if input_files is not None and isinstance(input_files, dict) and len(input_files) > 0: - self.warning("Using user provided input file(s) instead of the pre-defined file(s)") - urls = input_files - else: - urls = self.get_urls() - if urls is None: - raise ValueError("No input files provided") + urls = self.get_input_urls(input_files) self.info("Getting file(s) if not cached yet") paths = self.download_files(urls, cache) diff --git a/vecorel_cli/conversion/duckdb.py b/vecorel_cli/conversion/duckdb.py index 806a04d..74fc42b 100644 --- a/vecorel_cli/conversion/duckdb.py +++ b/vecorel_cli/conversion/duckdb.py @@ -116,13 +116,7 @@ def convert( if directory: os.makedirs(directory, exist_ok=True) - if input_files is not None and isinstance(input_files, dict) and len(input_files) > 0: - self.warning("Using user provided input file(s) instead of the pre-defined file(s)") - urls = input_files - else: - urls = self.get_urls() - if urls is None: - raise ValueError("No input files provided") + urls = self.get_input_urls(input_files) self.info("Getting file(s) if not cached yet") # Without a cache DuckDB reads remote parquet files in place; anything else, such as @@ -368,6 +362,8 @@ def write_query( compression = compression or "zstd" if compression == "zstd" and compression_level is None: compression_level = 15 + # DuckDB calls "none" "uncompressed"; the postprocess (Arrow) takes the original name + copy_compression = "uncompressed" if compression == "none" else compression geoparquet_version = geoparquet_version or "1.1.0" row_group_size = GeoParquet.row_group_size if isinstance(output_file, Path): @@ -561,12 +557,12 @@ def in_collections(cids): }} ) """, - [output_file, compression, collection_json], + [output_file, copy_compression, collection_json], ) if "id" in selected_targets and suffix_duplicate_ids: self._suffix_duplicate_ids_in_file( - con, output_file, compression, collection_json, row_group_size, per_collection + con, output_file, copy_compression, collection_json, row_group_size, per_collection ) # Sort against the same CRS-derived Hilbert grid as the @@ -591,7 +587,7 @@ def in_collections(cids): con, output_file, keys_path, - compression, + copy_compression, collection_json, final_group_size, ) diff --git a/vecorel_cli/conversion/per_file.py b/vecorel_cli/conversion/per_file.py new file mode 100644 index 0000000..1289fd7 --- /dev/null +++ b/vecorel_cli/conversion/per_file.py @@ -0,0 +1,53 @@ +import os +from tempfile import TemporaryDirectory + +from .duckdb import DuckDBBaseConverter + + +class PerFileConverterMixin: + """Converts each source on its own and merges the parts, for sources that do not fit + into memory together. The ids must be unique across the sources, and the parts need + about the size of the output in free disk space next to it.""" + + _part_urls = None + + def get_input_urls(self, input_files=None): + if self._part_urls is not None: + return self._part_urls + return super().get_input_urls(input_files) + + def convert(self, output_file, cache=None, input_files=None, variant=None, **kwargs) -> str: + self.select_variant(variant) + urls = self.get_input_urls(input_files) + sources = self._split_sources(urls) + # the variant is passed on as selected, so that its default is not reported per part + try: + if len(sources) <= 1: + self._part_urls = urls + return super().convert(output_file, cache=cache, variant=self.variant, **kwargs) + + directory = os.path.dirname(output_file) or "." + os.makedirs(directory, exist_ok=True) + with TemporaryDirectory(dir=directory) as parts_folder: + parts = [] + for index, source in enumerate(sources): + self.info(f"Converting source {index + 1}/{len(sources)}: {next(iter(source))}") + part = os.path.join(parts_folder, f"part_{index}.parquet") + self._part_urls = source + super().convert(part, cache=cache, variant=self.variant, **kwargs) + parts.append(part) + + self.info(f"Merging {len(parts)} parts into {output_file}") + merge_options = ("compression", "compression_level", "geoparquet_version") + return DuckDBBaseConverter().merge_parquet( + parts, output_file, **{k: kwargs[k] for k in merge_options if k in kwargs} + ) + finally: + self._part_urls = None + + def _split_sources(self, urls) -> list: + if isinstance(urls, str): + return [urls] + # the volumes of a multi-volume 7z archive can only be extracted together + regular, volume_groups = self._group_multivolume_7z(urls) + return [{uri: target} for uri, target in regular.items()] + list(volume_groups.values())