fix #24/import netcdf

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",

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)

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