Add typed Polars quantity frames

README.md

@@ -20,6 +20,8 @@ Every physical quantity in `encomp` carries a magnitude, a unit, and a dimension
 
 - **Static unit checking.** Common dimensionalities are encoded as `Quantity.__new__` overloads, and their `*` / `/` / `**` combinations as arithmetic-operator overloads, so a type checker flags a `Temperature` passed where a `Power` is expected. `@typeguard.typechecked` extends the same checks to runtime.
 
+- **Unit-carrying Polars columns** (`encomp.polars`). Declare `pressure = unit("bar")` on a `QuantityFrame`; validated attributes are typed `Quantity[..., pl.Expr]` values, and the physical units round-trip through Parquet/IPC as Arrow field metadata.
+
 - **`Fluid` / `Water` / `HumidAir`** (`encomp.fluids`) wrap [CoolProp](http://www.coolprop.org): fix a state with two points (three for humid air) and read any property as an attribute. Inputs and outputs are `Quantity` objects, so units are converted and validated automatically.
 
 - **CoolProp as Polars expressions** (`encomp.coolprop`). Properties evaluate as native Polars expression plugins written in Rust, so independent properties in one `select` / `with_columns` run in parallel without holding the GIL. See [Parallel CoolProp evaluation with Polars](#parallel-coolprop-evaluation-with-polars).
@@ -211,6 +213,85 @@ assert q1 == q2
 assert type(q1) is type(q2)
 ```
 
+## Unit-carrying Polars columns
+
+`Quantity` is the computation layer, while `UnitDType` is the DataFrame and storage layer. A `QuantityFrame` declares each unit once, validates the metadata restored from storage, and exposes typed quantity expressions. The attribute name is the Polars column name by default.
+
+```python
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import polars as pl
+
+from encomp.polars import QuantityFrame, unit, units_of
+from encomp.units import Unit
+
+
+class Sensors(QuantityFrame):
+    pressure = unit("bar")
+    flow = unit("m³/h")
+
+
+class Report(QuantityFrame):
+    power = unit("kW")
+
+
+raw = pl.DataFrame({"pressure": [1.0, 2.0], "flow": [10.0, 20.0]})
+stored = Sensors.from_untyped(raw)
+
+with TemporaryDirectory() as directory:
+    path = Path(directory) / "sensors.parquet"
+    stored.lf.sink_parquet(path)
+
+    sensors = Sensors.scan_parquet(path)
+    power = (sensors.pressure * sensors.flow).to("kW")
+    report = Report.derive(sensors, Report.power.assign(power))
+
+    assert units_of(report.lf)["power"] == Unit("kW")
+    assert abs(report.lf.collect()["power"].ext.storage()[0] - 0.2778) < 0.0001
+```
+
+The boundary is explicit in both directions:
+
+1. `Sensors.from_untyped(raw)` is the conspicuous assertion that bare input magnitudes have the declared units.
+2. `Sensors.scan_parquet(...)` validates metadata before typed attribute access becomes available. Compatible stored units are converted to the declaration inside the lazy plan.
+3. `Report.power.assign(power)` checks dimensionality, and `Report.derive(...)` converts to `kW` and restores its metadata without collecting.
+
+Use `unit("bar", name="Suction pressure")` when an external column name cannot be the Python attribute. Registered unit literals infer dimensionality for static checkers. Pint still accepts its open unit grammar: an unlisted spelling such as `unit("kg/(m*s**2)")` is inferred at runtime and typed conservatively as unknown; add `asdim=Pressure` when precise static typing is required. The override is validated rather than cast.
+
+The validated expressions pass directly into the high-level fluid API. Non-SI declarations are converted before the native plugin runs, and the result can be persisted with another declaration:
+
+```python
+import polars as pl
+
+from encomp.fluids import Water
+from encomp.polars import QuantityFrame, unit, units_of
+from encomp.units import Unit
+
+
+class States(QuantityFrame):
+    pressure = unit("bar", name="P")
+    temperature = unit("degC", name="T")
+
+
+class Properties(QuantityFrame):
+    density = unit("kg/m³", name="rho")
+
+
+states = States.from_untyped(pl.LazyFrame({"P": [5.0], "T": [150.0]}))
+water = Water[pl.Expr](P=states.pressure, T=states.temperature)
+properties = Properties.derive(states, Properties.density.assign(water.D))
+
+assert units_of(properties.lf)["rho"] == Unit("kg/m³")
+assert properties.lf.collect()["rho"].ext.storage()[0] > 900.0
+```
+
+Raw arithmetic on a unit-typed column is refused because Polars cannot delegate third-party unit algebra. This prevents accidental operations such as adding pressure to flow; intentional computation goes through `Quantity`. Syntax normalization is conservative: `m^3` and `m³` produce the same dtype, while dimensionally equivalent renderings such as `Pa` and `N/m²` remain different metadata identities.
+
+The persisted Arrow field contract is `ARROW:extension:name = "encomp.unit"` and `ARROW:extension:metadata = <canonical unit>`. Import `encomp.polars`—or `encomp.coolprop`, which registers the same dtype—before reading with Polars 1.x. For readers that do not import encomp, `POLARS_UNKNOWN_EXTENSION_TYPE_BEHAVIOR=load_as_extension` preserves generic extension metadata; otherwise current Polars 1.x warns and loads the numeric storage type.
+
+The Polars extension API is marked unstable upstream, so this integration is experimental even though encomp pins its required behavior in tests.
+
 ## The `Fluid` class
 
 `encomp.fluids.Fluid` wraps the *CoolProp* library.

docs/api.md

@@ -26,6 +26,14 @@ Autogenerated reference for the documented public `encomp` modules. The core mod
    :show-inheritance:
 ```
 
+## encomp.polars
+
+```{eval-rst}
+.. automodule:: encomp.polars
+   :members:
+   :show-inheritance:
+```
+
 ## encomp.coolprop
 
 ```{eval-rst}

docs/usage.md

@@ -305,6 +305,91 @@ type(Q(pl.lit(5), "kg").m)  # pl.Expr
 
 A `Quantity` with a `pl.Expr` magnitude is a deferred plan, not data. Only unit algebra (arithmetic, comparison, `.to`, `abs`) is meaningful on it; reach the underlying Polars object with `.m` to compute inside a `select` / `with_columns`. This is also how the parallel CoolProp evaluation described below is driven.
 
+### Persisting units in Polars schemas
+
+A bare `Quantity(pl.col("P"), "bar")` carries its unit only in the Python wrapper. The association disappears if another service reads the underlying Parquet file. `encomp.polars` supplies the missing frame-level layer: `UnitDType` stores the unit in Arrow field metadata, while the numeric values remain the extension dtype's storage.
+
+Declare the producer contract once on a {py:class}`encomp.polars.QuantityFrame`. Constructing it normally validates units already present in the Polars schema; {py:meth}`encomp.polars.QuantityFrame.from_untyped` is the explicit operation that assigns declared units to bare input data:
+
+```python
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import polars as pl
+
+from encomp.polars import QuantityFrame, unit, units_of
+from encomp.units import Unit
+
+
+class Sensors(QuantityFrame):
+    pressure = unit("bar", name="P")
+    flow = unit("m³/h", name="V")
+
+
+class Report(QuantityFrame):
+    power = unit("kW", name="Hydraulic power")
+
+
+source = Sensors.from_untyped(pl.DataFrame({"P": [1.0, 2.0], "V": [10.0, 20.0]}))
+
+with TemporaryDirectory() as directory:
+    path = Path(directory) / "source.parquet"
+    source.lf.sink_parquet(path)
+    sensors = Sensors.scan_parquet(path)
+
+    power = sensors.pressure * sensors.flow
+    result = Report.derive(sensors, Report.power.assign(power))
+
+    assert units_of(result.lf) == {
+        "P": Unit("bar"),
+        "V": Unit("m³/h"),
+        "Hydraulic power": Unit("kW"),
+    }
+    assert result.lf.collect()["Hydraulic power"].ext.storage().to_list() == [
+        0.2777777777777778,
+        1.1111111111111112,
+    ]
+```
+
+The declaration `pressure = unit("bar", name="P")` contains the physical information once. `name=` is only needed because the external name differs from the attribute; `pressure = unit("bar")` would target a column named `"pressure"`. Known unit literals make `sensors.pressure` statically `Quantity[Pressure, pl.Expr]`. Class access returns the typed target declaration instead, so `Report.power.assign(...)` rejects a pressure quantity under all three supported type checkers. The lazy magnitude is still an ordinary `pl.col("P").ext.storage()` expression, so Polars can optimize it normally.
+
+Unit strings are not restricted to the literal autocomplete list. Pint parses other valid expressions and encomp infers their dimensionality at runtime. Since a static checker cannot execute Pint's registry, an unlisted string produces `Column[UnknownDimensionality]`; supply the exceptional `asdim=Pressure` argument when that column needs precise static typing. Compatibility is checked immediately.
+
+Construction safely converts compatible stored units to the declaration inside the lazy plan; incompatible dimensionalities fail before data is read. Schema classes compose through ordinary inheritance when one validated view needs columns from multiple schemas.
+
+The same bridge composes with the high-level fluid API without dropping to raw expressions. `Water` accepts the descriptor quantities, performs its required SI conversions, and returns a typed expression quantity that can be assigned straight back to a declared output:
+
+```python
+import polars as pl
+
+from encomp.fluids import Water
+from encomp.polars import QuantityFrame, unit, units_of
+from encomp.units import Unit
+
+
+class States(QuantityFrame):
+    pressure = unit("bar", name="P")
+    temperature = unit("degC", name="T")
+
+
+class Properties(QuantityFrame):
+    density = unit("kg/m³", name="rho")
+
+
+states = States.from_untyped(pl.LazyFrame({"P": [5.0], "T": [150.0]}))
+water = Water[pl.Expr](P=states.pressure, T=states.temperature)
+properties = Properties.derive(states, Properties.density.assign(water.D))
+
+assert units_of(properties.lf)["rho"] == Unit("kg/m³")
+assert properties.lf.collect()["rho"].ext.storage()[0] > 900.0
+```
+
+This convenience belongs to `Water`/`Fluid`, where units are present in the `Quantity` inputs. The lower-level `encomp.coolprop.water()` function accepts bare Polars columns and therefore cannot insert unit conversions: extension-typed inputs must already carry the exact SI unit, or callers must pass an explicitly converted magnitude such as `states.pressure.to("Pa").m`.
+
+Polars has no extension hook for encomp's multiplication or supertype rules. It therefore refuses value-producing operations directly on unit-typed columns; computation must explicitly cross into `Quantity`. Value-preserving operations such as filtering, sorting, aliases, join keys and group keys keep the dtype, and incompatible unit dtypes cannot be concatenated.
+
+Parquet and IPC store `ARROW:extension:name = "encomp.unit"` and `ARROW:extension:metadata = <canonical unit>` on each field. The canonical rendering normalizes syntax, not dimensional equivalence: `m^3` and `m³` match, but `Pa` and `N/m²` remain distinct dtype identities. The upstream Polars extension API is unstable, so encomp treats this integration as experimental and pins its required I/O and refusal behavior in tests.
+
 ### Combining quantities
 
 The output of an operation on quantities is always consistent with the input dimensionalities.
@@ -313,6 +398,8 @@ Inconsistent or ambiguous operations raise descriptive errors.
 Units do not always cancel out automatically.
 Call {py:meth}`encomp.units.Quantity.to_base_units` to simplify to base SI units, {py:meth}`encomp.units.Quantity.to` when the target unit is known, or {py:meth}`encomp.units.Quantity.to_reduced_units` to cancel units without converting to base SI units.
 
+When only the converted magnitude is needed, {py:meth}`encomp.units.Quantity.m_as` is the typed shorthand for `.to(unit).m`; it preserves the magnitude container type and applies the same dimensionality and temperature checks.
+
 ```python
 from encomp.units import Quantity as Q
 

encomp/_polars_dtype.py

@@ -0,0 +1,79 @@
+"""Pint-light registration core for encomp's Polars unit extension dtype.
+
+This module may be imported by :mod:`encomp.coolprop` before any unit computation.
+Keep Pint and :mod:`encomp.units` imports inside the methods that actually need them:
+registering the dtype itself must be cheap and must happen before Polars reads a file.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+
+if TYPE_CHECKING:
+    from .units import Unit
+
+EXTENSION_NAME = "encomp.unit"
+"""Public, stable Arrow extension name for a unit-carrying column."""
+
+_CANONICAL_UNIT_FORMAT = "~P"
+"""Fixed Pint format spec for the unit string stored in dtype metadata."""
+
+
+def canonical_unit_string(unit: str | Unit[Any]) -> str:
+    """Return the fixed canonical rendering stored in extension metadata.
+
+    The rendering is independent of encomp's process-wide display format because it
+    is an on-disk, cross-process contract. It normalizes spelling (``"m^3"`` and
+    ``"m³"``), not physical equivalence (``"Pa"`` and ``"N/m²"`` remain distinct).
+    """
+    from .units import Unit
+
+    parsed = Unit(unit) if isinstance(unit, str) else unit
+    return format(parsed, _CANONICAL_UNIT_FORMAT)
+
+
+def _validate_storage(storage: pl.DataType) -> None:
+    if not (storage.is_float() or storage.is_integer()):
+        raise TypeError(
+            f"unit dtype storage must be a float or integer dtype, got {storage!r}. "
+            "Boolean and non-numeric columns cannot carry a unit, and a column that "
+            "already has a unit dtype cannot be re-wrapped."
+        )
+
+
+class UnitDType(pl.BaseExtension):
+    """Polars extension dtype carrying a physical unit as column metadata.
+
+    A dtype instance consists of the stable extension name ``"encomp.unit"``, a
+    numeric storage dtype, and a canonical unit string. Polars refuses value-producing
+    arithmetic on extension columns; validated :class:`encomp.polars.QuantityFrame`
+    descriptors enter Quantity unit algebra and ``QuantityFrame.derive`` returns
+    results to a frame.
+    """
+
+    def __init__(self, unit: str | Unit[Any], storage: pl.DataType | None = None) -> None:
+        if storage is None:
+            storage = pl.Float64()
+        _validate_storage(storage)
+        super().__init__(EXTENSION_NAME, storage, canonical_unit_string(unit))
+
+    @property
+    def unit(self) -> Unit[Any]:
+        """The column unit parsed from the extension metadata."""
+        from .units import Unit
+
+        metadata = self.ext_metadata()
+        if metadata is None:
+            raise ValueError(f"{EXTENSION_NAME} dtype without a unit string in its metadata: {self!r}")
+        return Unit(metadata)
+
+    def _string_repr(self) -> str:
+        return f"unit[{self.ext_metadata()}]"
+
+
+# This import side effect is the purpose of the module. Both encomp.polars and every
+# encomp entry point that can feed extension-typed data to the native plugin import it
+# before reading data, so Polars never strips the unit metadata first.
+pl.register_extension_type(EXTENSION_NAME, UnitDType)

encomp/coolprop/__init__.py

@@ -32,12 +32,14 @@
 import warnings
 from functools import cache, lru_cache
 from pathlib import Path
-from typing import Any, Literal, TypeIs, get_args
+from typing import Any, Literal, TypeIs, cast, get_args
 
 import CoolProp.CoolProp as _CoolProp
 import polars as pl
 from polars.plugins import register_plugin_function
 
+from .._polars_dtype import canonical_unit_string as _canonical_unit_string
+
 # CoolProp ships incomplete type stubs; treat the module as Any (mirrors encomp).
 _cp: Any = _CoolProp
 
@@ -502,6 +504,32 @@ def _input_name(x: str | pl.Expr) -> str:
     return x if isinstance(x, str) else x.meta.output_name()
 
 
+@cache
+def _expected_si_unit(name: str) -> str:
+    """Canonical rendering of the SI unit CoolProp expects for state input ``name``.
+
+    Derived from CoolProp's own parameter metadata (no unit table of our own) and
+    rendered by the same internal canonicalization the unit dtype applies to its
+    metadata, so the plugin can accept a unit-typed
+    input by comparing the two strings for equality.
+    """
+    return _canonical_unit_string(_cp.get_parameter_information(_cp.get_parameter_index(name), "units"))
+
+
+@cache
+def _expected_ha_si_unit(name: str) -> str:
+    """Canonical SI unit for HAPropsSI input ``name`` (see :func:`_expected_si_unit`).
+
+    HAPropsSI has no parameter-information API, so this reads
+    :meth:`encomp.fluids.HumidAir.get_coolprop_unit` -- the existing single source for
+    humid-air units (a lookup failure means that mapping has a gap and raises here).
+    """
+    from ..fluids import CProperty, HumidAir  # deferred: encomp.fluids imports this module at top level
+
+    # runtime-validated dynamic lookup: get_coolprop_unit raises for an unknown name
+    return _canonical_unit_string(HumidAir.get_coolprop_unit(cast(CProperty, name)))
+
+
 def _reject_duplicate_input_keys(kind: str, keyed_names: tuple[tuple[str, int | str], ...]) -> None:
     seen: dict[int | str, str] = {}
 
@@ -632,6 +660,13 @@ def fluid(
     column read from it; an expression by its output name (``pl.col("P")`` or
     ``pl.col("p").alias("P")``). Both must be CoolProp state inputs (P/T/Q/D/H/S/U,
     any pair: PT, PH, PQ, ...).
+
+    Inputs are numeric columns holding SI magnitudes. A unit-carrying extension dtype
+    (:mod:`encomp.polars`) is accepted when its unit is already the SI unit CoolProp
+    expects for that input (verified against CoolProp's own parameter metadata); any
+    other unit is refused at schema-resolution time with the expected unit named --
+    convert first (e.g. ``encomp.polars.quantities`` + ``.to(...)``), or pass raw SI
+    values via ``.ext.storage()``.
     """
     name1, name2 = _input_name(input1), _input_name(input2)
     for nm in (name1, name2):
@@ -647,6 +682,7 @@ def fluid(
     )
     pair_idx, swap = _resolve_pair(name1, name2)
     a, b = (input2, input1) if swap else (input1, input2)  # canonical order
+    name_a, name_b = (name2, name1) if swap else (name1, name2)
     a_expr, b_expr = _as_expr(a), _as_expr(b)
     backend, fluids, fractions = resolve_fluid_spec(name, composition)
     phase = _phase_from_assumed(assume_phase) if assume_phase is not None else None
@@ -677,6 +713,7 @@ def fluid(
             "phase": phase,
             "fractions": fractions,
             "scalar_mask": [_is_scalar(a_expr), _is_scalar(b_expr)],
+            "expected_units": [_expected_si_unit(name_a), _expected_si_unit(name_b)],
         },
         is_elementwise=True,
         use_abs_path=True,
@@ -711,6 +748,9 @@ def humid_air(
 
     Each input identifies its property by name (the string, or the expression's
     output name); all three must be HAPropsSI state inputs (T, P, R, W, B, ...).
+    Like :func:`fluid`, inputs hold SI magnitudes: a unit-carrying extension dtype is
+    accepted only when its unit is already the expected SI unit, and refused with the
+    expected unit named otherwise.
     """
     # Unlike the fluid path (validated by CoolProp's param_index in the plugin), HAPropsSI
     # returns _HUGE -> null for an unknown output, so a typo'd output would silently yield
@@ -745,6 +785,7 @@ def humid_air(
             "name2": name2,
             "name3": name3,
             "scalar_mask": [_is_scalar(e1), _is_scalar(e2), _is_scalar(e3)],
+            "expected_units": [_expected_ha_si_unit(name1), _expected_ha_si_unit(name2), _expected_ha_si_unit(name3)],
         },
         is_elementwise=True,
         use_abs_path=True,