diff --git a/pyproject.toml b/pyproject.toml index 205edc2..a547cd9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,6 +106,9 @@ cuda = [ test = [ "damast[ml]", "coverage", + # NetCDF loading (see PolarsDataFrame.import_netcdf) + "netCDF4", + "xarray", "pandas>=2", "pytest", "pytest-console-scripts", diff --git a/src/damast/core/dataframe.py b/src/damast/core/dataframe.py index ca256ee..b8502d9 100644 --- a/src/damast/core/dataframe.py +++ b/src/damast/core/dataframe.py @@ -402,7 +402,7 @@ def from_files(cls, return cls(dataframe=df, metadata=metadata, validation_mode=validation_mode) @classmethod - def load_parquet(cls, files) -> tuple[AnnotatedDataFrame, dict[str, MetaData]]: + def load_parquet(cls, files) -> tuple[polars.LazyFrame, dict[str, MetaData]]: _log.info(f"Loading parquet: {files=}") metadata_per_file = {} @@ -427,17 +427,17 @@ def load_parquet(cls, files) -> tuple[AnnotatedDataFrame, dict[str, MetaData]]: return df, metadata_per_file @classmethod - def load_netcdf(cls, files) -> tuple[AnnotatedDataFrame, dict[str, MetaData]]: + def load_netcdf(cls, files) -> tuple[polars.LazyFrame, dict[str, MetaData]]: _log.info(f"Loading netcdf: {files=}") return XDataFrame.import_netcdf(files) @classmethod - def load_hdf(cls, files) -> tuple[AnnotatedDataFrame, dict[str, MetaData]]: + def load_hdf(cls, files) -> tuple[polars.LazyFrame, dict[str, MetaData]]: _log.info(f"Loading hdf: {files=}") return XDataFrame.import_hdf5(files) @classmethod - def load_csv(cls, files) -> tuple[AnnotatedDataFrame, dict[str, MetaData]]: + def load_csv(cls, files) -> tuple[polars.LazyFrame, dict[str, MetaData]]: _log.info(f"Loading csv: {files=}") df = polars.scan_csv(files, separator=";", **DAMAST_CSV_DEFAULT_ARGS) diff --git a/src/damast/core/polars_dataframe.py b/src/damast/core/polars_dataframe.py index c941e0c..7946664 100644 --- a/src/damast/core/polars_dataframe.py +++ b/src/damast/core/polars_dataframe.py @@ -2,15 +2,17 @@ import ast import logging +import math import os import re from pathlib import Path -from typing import ClassVar +from typing import Any, ClassVar import numpy as np import polars import polars.api from polars import LazyFrame +from polars.io.plugins import register_io_source from pydantic import ValidationError from damast.utils import ensure_packages @@ -546,24 +548,184 @@ def from_vaex_hdf5(cls, path: str | Path) -> tuple[polars.LazyFrame, 'MetaData'] return polars.LazyFrame(data), metadata + #: Maximum number of grid cells read per batch by the lazy NetCDF scan (before dropping padding) + NETCDF_BATCH_SIZE: ClassVar[int] = 100_000 + @classmethod def import_netcdf(cls, path: list[str|Path]) -> tuple[polars.LazyFrame, dict[str, 'MetaData']]: #noqa + """ + Lazily scan NetCDF files - see :func:`scan_netcdf` - and extract metadata from their CF + attributes, see :func:`_metadata_from_cf_attributes`. + """ + frames = [] + metadata = {} + for f in path: + lazyframe, variables = cls.scan_netcdf(f) + frames.append(lazyframe) - ensure_packages(pkgs=["dask", "xarray", "pandas"], - required_for="Loading netcdf files", - hints=", additionally either netcdf4 or h5netcdf have to be installed") + file_metadata = cls._metadata_from_cf_attributes(lazyframe.collect_schema(), variables, + source=Path(f).name) + if file_metadata is not None: + metadata[str(f)] = file_metadata + + return polars.concat(frames, how="diagonal_relaxed"), metadata - import pandas as pd + @classmethod + def scan_netcdf(cls, path: str | Path) -> tuple[polars.LazyFrame, dict[str, tuple[dict, dict]]]: + """ + Lazily scan a NetCDF file as a table with one row per grid cell - the same layout as + ``xarray.Dataset.to_dataframe()``, with the dimensions as leading columns. + + Nothing is read until the frame is collected. The grid is then read in slices along its + first dimension, so memory is bounded by a slice rather than the whole grid, and rows are + filtered/projected/limited per slice. Rows in which every data variable spanning the full + grid is missing - e.g. the padding of a sparse (entity x time) grid - are dropped. + + :param path: The NetCDF file + :return: The lazyframe, and variable name -> (CF attributes, xarray encoding) + """ + ensure_packages(pkgs=["xarray"], + required_for="Loading netcdf files", + hint="additionally either netCDF4 or h5netcdf have to be installed") import xarray - dataframes = [] - for f in path: - ds = xarray.open_dataset(f) - dataframes.append( ds.to_dataframe().reset_index() ) - pandas_df = pd.concat(dataframes, ignore_index=True).reset_index() - df = polars.from_pandas(pandas_df) + with xarray.open_dataset(path) as ds: + variables = {name: (dict(variable.attrs), dict(variable.encoding)) + for name, variable in ds.variables.items()} + schema = cls._netcdf_schema(ds) + + def read_batches(with_columns: list[str] | None, + predicate: polars.Expr | None, + n_rows: int | None, + batch_size: int | None): + with xarray.open_dataset(path) as ds: + dims = list(ds.sizes) + # A cell is padding if all variables spanning the full grid are missing there - lower + # dimensional ones (e.g. static per-entity values) are just repeated into every cell + data_vars = [name for name, var in ds.data_vars.items() if set(var.dims) == set(dims)] + data_vars = data_vars or list(ds.data_vars) + # cells per step along the first dimension - which to_dataframe() iterates slowest + cells_per_step = math.prod(list(ds.sizes.values())[1:]) + # polars' batch_size is only a hint - cap it, so memory stays bounded per slice + max_cells = min(batch_size or cls.NETCDF_BATCH_SIZE, cls.NETCDF_BATCH_SIZE) + step = max(1, max_cells // max(1, cells_per_step)) + first_dim_size = ds.sizes[dims[0]] if dims else 1 + + for start in range(0, first_dim_size, step): + if n_rows is not None and n_rows <= 0: + return + + part = ds.isel({dims[0]: slice(start, start + step)}) if dims else ds + pandas_df = part.to_dataframe().reset_index() + if data_vars: + pandas_df = pandas_df.dropna(how="all", subset=data_vars) + + # e.g. an all-missing string column would otherwise come back as Null + df = polars.from_pandas(pandas_df).cast(schema) + if predicate is not None: + df = df.filter(predicate) + if with_columns is not None: + df = df.select(with_columns) + if n_rows is not None: + df = df.head(n_rows) + n_rows -= df.height + yield df + + return register_io_source(read_batches, schema=schema), variables + + @staticmethod + def _netcdf_schema(ds) -> polars.Schema: + """ + Columns and dtypes of ``ds.to_dataframe()`` without reading data: taken from an empty + slice, where object columns (strings) cannot be inferred and default to String. + """ + empty = ds.isel({dim: slice(0, 0) for dim in ds.sizes}).to_dataframe().reset_index() + return polars.Schema({ + column: polars.String if dtype.kind == "O" else polars.from_pandas(empty[column]).dtype + for column, dtype in empty.dtypes.items() + }) + + @classmethod + def _metadata_from_cf_attributes(cls, + schema: polars.Schema, + variables: dict[str, tuple[dict, dict]], + source: str) -> 'MetaData' | None: # noqa + """ + Create metadata for the columns of a loaded NetCDF file from the CF attributes of its + variables: 'long_name' becomes the description, 'units' the unit - if it can be parsed -, + and 'valid_range'/'valid_min'/'valid_max' the value range of a numeric column. + + '_FillValue'/'missing_value' are not mapped: xarray already decodes them to NaN (null in + polars), while damast's missing_value is the value used to replace out-of-range values. + + :param variables: variable name -> (attributes, xarray encoding) + :return: The metadata, or None if no variable carries any of these attributes - so that + callers can fall back to searching for a spec file or inferring the metadata + """ + # avoid circular dependencies + from damast.core.annotations import Annotation + from damast.core.data_description import MinMax + from damast.core.metadata import DataSpecification, MetaData + from damast.core.units import Unit + + column_specs = [] + has_cf_attributes = False + for column, dtype in schema.items(): + attrs, encoding = variables.get(column, ({}, {})) + spec = DataSpecification(name=column, representation_type=dtype) + + if "long_name" in attrs: + spec.description = str(attrs["long_name"]) + has_cf_attributes = True + + if "units" in attrs: + has_cf_attributes = True + try: + spec.unit = Unit(str(attrs["units"])) + except ValueError: + logger.info(f"NetCDF {source}: cannot interpret unit '{attrs['units']}' of '{column}' - ignoring it") + + valid_range = cls._cf_valid_range(attrs, encoding) + if valid_range is not None: + has_cf_attributes = True + # e.g. a decoded time column cannot be compared with its (numeric) raw range + if dtype.is_numeric(): + spec.value_range = MinMax(*valid_range) + else: + logger.info(f"NetCDF {source}: ignoring valid range of non-numeric '{column}'") + + column_specs.append(spec) + + if not has_cf_attributes: + return None + + return MetaData(columns=column_specs, + annotations=[Annotation(name=Annotation.Key.Source, value=source)]) + + @staticmethod + def _cf_valid_range(attrs: dict, encoding: dict) -> tuple[Any, Any] | None: + """ + (min, max) from the CF 'valid_range', or 'valid_min'/'valid_max' attributes - an open side + becomes -inf/inf. CF defines them in packed units, so they are unpacked like the data via + 'scale_factor'/'add_offset', which xarray moves into the variable's encoding. + + :return: The range, or None if the variable declares none + """ + if "valid_range" in attrs: + low, high = np.asarray(attrs["valid_range"]).tolist() + elif "valid_min" in attrs or "valid_max" in attrs: + low = np.asarray(attrs.get("valid_min", -np.inf)).item() + high = np.asarray(attrs.get("valid_max", np.inf)).item() + else: + return None + + if "scale_factor" in encoding or "add_offset" in encoding: + scale = float(encoding.get("scale_factor", 1.0)) + offset = float(encoding.get("add_offset", 0.0)) + # a negative scale_factor swaps the bounds + low, high = sorted([low * scale + offset, high * scale + offset]) - return df.lazy(), {} + return low, high @classmethod def import_hdf5(cls, files: str | Path | list[str|Path]) -> tuple[polars.LazyFrame, dict[str, 'MetaData']]: # noqa diff --git a/tests/damast/core/test_dataframe.py b/tests/damast/core/test_dataframe.py index a357dcf..a5ecfe6 100644 --- a/tests/damast/core/test_dataframe.py +++ b/tests/damast/core/test_dataframe.py @@ -9,6 +9,7 @@ from astropy import units from damast.core.annotations import Annotation +from damast.core.constants import DAMAST_SUPPORTED_FILE_FORMATS from damast.core.data_description import ListOfValues, MinMax from damast.core.dataframe import AnnotatedDataFrame from damast.core.metadata import ( @@ -495,3 +496,208 @@ def test_update_preserves_representation_type_when_step_output_declares_none(): assert adf.metadata["date_time_utc"].description == "original description" assert adf.metadata["date_time_utc"].unit == units.deg adf.validate_metadata() + + +def _write_parquet(df: polars.DataFrame, path: Path): + df.write_parquet(path) + + +def _write_csv(df: polars.DataFrame, path: Path): + df.write_csv(path, separator=";") + + +def _write_hdf(df: polars.DataFrame, path: Path): + # as AnnotatedDataFrame.save does: the loader needs the per-column metadata nodes + XDataFrame.export_hdf5(df, path) + AnnotatedDataFrame.infer_annotation(df).append_to_hdf(path) + + +def _write_netcdf(df: polars.DataFrame, path: Path): + import xarray + + xarray.Dataset.from_dataframe(df.to_pandas()).to_netcdf(path) + + +ROUND_TRIP_WRITERS = { + "parquet": _write_parquet, + "csv": _write_csv, + "hdf": _write_hdf, + "netcdf": _write_netcdf, +} + + +@pytest.mark.parametrize(["file_format", "suffix"], [ + [file_format, suffix] + for file_format, suffixes in DAMAST_SUPPORTED_FILE_FORMATS.items() + for suffix in suffixes +]) +def test_from_files_round_trip_for_every_supported_suffix(file_format, suffix, tmp_path): + """Every registered suffix must reach a working loader - including its optional-package checks.""" + expected = polars.DataFrame({"x": [1, 2, 3], "y": [0.5, 1.5, 2.5]}) + path = tmp_path / f"data{suffix}" + ROUND_TRIP_WRITERS[file_format](expected, path) + + adf = AnnotatedDataFrame.from_files(files=[str(path)], metadata_required=False) + + polars.testing.assert_frame_equal(adf.lazyframe.select(["x", "y"]).collect(), expected) + + +def test_from_files_netcdf_uses_cf_attributes_as_metadata(tmp_path): + import xarray + + from damast.core.units import Unit + + path = tmp_path / "data.nc" + xarray.Dataset( + { + "speed": ("row", [1.0, 2.0, 3.0], {"units": "m s-1", "long_name": "Speed over ground"}), + "lat": ("row", [60.0, 61.0, 62.0], {"units": "degrees_north", "long_name": "Latitude"}), + }, + coords={"row": [0, 1, 2]}, + ).to_netcdf(path) + + # metadata is required by default - it has to come from the file itself here + adf = AnnotatedDataFrame.from_files(files=[str(path)]) + + assert not adf.metadata_inferred + assert adf.metadata["speed"].unit == Unit("m s-1") + assert adf.metadata["speed"].description == "Speed over ground" + assert adf.metadata["speed"].representation_type == polars.Float64 + # 'degrees_north' is not a parseable unit, but the description is still used + assert adf.metadata["lat"].unit is None + assert adf.metadata["lat"].description == "Latitude" + + +def test_from_files_netcdf_without_cf_attributes_falls_back_to_inference(tmp_path): + import xarray + + path = tmp_path / "data.nc" + xarray.Dataset({"speed": ("row", [1.0, 2.0, 3.0])}).to_netcdf(path) + + adf = AnnotatedDataFrame.from_files(files=[str(path)], metadata_required=False) + + assert adf.metadata_inferred + + +def test_from_files_netcdf_uses_cf_valid_range_as_value_range(tmp_path): + import xarray + + path = tmp_path / "data.nc" + xarray.Dataset( + { + # packed: stored as int16, data = raw * 0.1 + 10 - valid_range is given in raw units + "speed": ("row", [10.0, 20.0, 30.0], {"valid_range": np.array([0, 400], dtype="int16")}), + "depth": ("row", [1.0, 2.0, 3.0], {"valid_min": 0.0}), + }, + coords={ + "row": [0, 1, 2], + "time": ("row", pd.date_range("2026-01-01", periods=3), {"valid_min": 0}), + }, + ).to_netcdf(path, encoding={"speed": {"dtype": "int16", "scale_factor": 0.1, "add_offset": 10.0, + "_FillValue": -32768}}) + + metadata = AnnotatedDataFrame.from_files(files=[str(path)]).metadata + + assert (metadata["speed"].value_range.min, metadata["speed"].value_range.max) == pytest.approx((10.0, 50.0)) + # an open side stays unbounded + assert (metadata["depth"].value_range.min, metadata["depth"].value_range.max) == (0.0, np.inf) + # a raw numeric range cannot apply to the decoded datetime column + assert metadata["time"].value_range is None + + +def _sparse_grid_dataset(n_mmsi: int, n_time: int, observed_ratio: float, seed: int = 0): + """An AIS-like (mmsi x time) grid, padded with NaN where a vessel has no observation.""" + import xarray + + rng = np.random.default_rng(seed) + observed = rng.uniform(size=(n_mmsi, n_time)) < observed_ratio + observed[0, 0] = True + speed = np.where(observed, rng.uniform(0, 20, (n_mmsi, n_time)), np.nan) + lat = np.where(observed, rng.uniform(55, 70, (n_mmsi, n_time)), np.nan) + return xarray.Dataset( + { + "speed": (("mmsi", "time"), speed), + "lat": (("mmsi", "time"), lat), + # static per vessel - repeated into every cell, so it must not count as an observation + "ship_type": ("mmsi", rng.integers(30, 90, n_mmsi)), + }, + coords={"mmsi": np.arange(n_mmsi) + 257_000_000, + "time": pd.date_range("2026-01-01", periods=n_time, freq="min")}, + ) + + +def _without_padding(path: Path) -> polars.DataFrame: + """What an eager read of the file gives, minus the padding cells.""" + import xarray + + with xarray.open_dataset(path) as ds: + pandas_df = ds.to_dataframe().reset_index() + return polars.from_pandas(pandas_df.dropna(how="all", subset=["speed", "lat"])) + + +@pytest.fixture +def small_netcdf_batches(monkeypatch): + # a single mmsi (one row of the grid) per batch, to exercise multi-batch reads + monkeypatch.setattr(XDataFrame, "NETCDF_BATCH_SIZE", 1) + + +def test_scan_netcdf_matches_eager_to_dataframe_without_padding(tmp_path, small_netcdf_batches): + ds = _sparse_grid_dataset(n_mmsi=20, n_time=30, observed_ratio=0.15) + path = tmp_path / "grid.nc" + ds.to_netcdf(path) + + df, _ = XDataFrame.import_netcdf([path]) + result = df.collect() + + polars.testing.assert_frame_equal(result, _without_padding(path)) + assert result.height == int(ds["speed"].notnull().sum()) + + +def test_scan_netcdf_concatenates_multiple_files(tmp_path, small_netcdf_batches): + datasets = [_sparse_grid_dataset(n_mmsi=5, n_time=8, observed_ratio=0.3, seed=seed) for seed in [1, 2]] + paths = [tmp_path / f"grid_{i}.nc" for i in range(len(datasets))] + for ds, path in zip(datasets, paths): + ds.to_netcdf(path) + + df, _ = XDataFrame.import_netcdf(paths) + + polars.testing.assert_frame_equal(df.collect(), polars.concat([_without_padding(path) for path in paths])) + + +def test_scan_netcdf_applies_filter_and_projection(tmp_path, small_netcdf_batches): + ds = _sparse_grid_dataset(n_mmsi=10, n_time=10, observed_ratio=0.5) + path = tmp_path / "grid.nc" + ds.to_netcdf(path) + + df, _ = XDataFrame.import_netcdf([path]) + result = df.filter(polars.col("speed") > 10).select(["mmsi", "speed"]).collect() + + expected = _without_padding(path).filter(polars.col("speed") > 10).select(["mmsi", "speed"]) + polars.testing.assert_frame_equal(result, expected) + + +def test_scan_netcdf_reads_nothing_until_collected_and_stops_early(tmp_path, monkeypatch, small_netcdf_batches): + import xarray + + ds = _sparse_grid_dataset(n_mmsi=50, n_time=10, observed_ratio=0.5) + path = tmp_path / "grid.nc" + ds.to_netcdf(path) + + rows_read = [] + to_dataframe = xarray.Dataset.to_dataframe + + def spy(self, *args, **kwargs): + result = to_dataframe(self, *args, **kwargs) + rows_read.append(len(result)) + return result + + monkeypatch.setattr(xarray.Dataset, "to_dataframe", spy) + + df, _ = XDataFrame.import_netcdf([path]) + lazy_query = df.filter(polars.col("speed") > 0).head(1) + # only the (empty) schema probe so far + assert sum(rows_read) == 0 + + assert lazy_query.collect().height == 1 + # far fewer than the 500 grid cells + assert 0 < sum(rows_read) < 50