Polars2

.gitignore

@@ -147,6 +147,8 @@ Thumbs.db
 # PyCharm
 /.idea
 *.iml
+.zed
+Test Results*.html
 
 # Dispatch
 *.zip

docs/conf.py

@@ -89,6 +89,7 @@
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None),
     "pandera": ("https://pandera.readthedocs.io/en/stable", None),
     "plotly": ("https://plotly.com/python-api-reference/", None),
+    "polars": ("https://pola-rs.github.io/polars/py-polars/html/reference/", None),
     "pytest": ("https://docs.pytest.org/en/latest/", None),
     "python": ("https://docs.python.org/3", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/", None),

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools<66", "setuptools_scm[toml]<8", "wheel"]
+requires = ["setuptools<69", "setuptools_scm[toml]<8", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -24,6 +24,7 @@ dependencies = [
     "numpy >= 1.18.5,<2",
     "pandas >= 1.4,< 2.1",
     "pandera >= 0.12, < 0.16",
+    "polars >=0.17, < 0.19",
     "pyarrow>=7, <13",
     "rmi.etoolbox @ git+https://github.com/rmi/etoolbox.git",
 ]

src/dispatch/metadata.py

@@ -4,6 +4,7 @@
 
 import pandas as pd
 import pandera as pa
+import polars as pl
 
 LOGGER = logging.getLogger(__name__)
 
@@ -255,3 +256,83 @@ def renewables(
                 "`re_profiles` and `load_profile` indexes must match"
             ) from exc
         return re_plant_specs, re_profiles
+
+
+class IDConverter:
+    """Helper for converting ids.
+
+    Converting between
+    :mod: `pandas` and
+    :mod: `polars`, especially for
+    multi-level columns.
+    """
+
+    __slots__ = (
+        "dt",
+        "disp_convert",
+        "disp_big_idx",
+        "re_convert",
+        "re_big_idx",
+        "storage_convert",
+        "storage_big_idx",
+    )
+
+    def __init__(  # noqa: D107
+        self, dispatchable_specs, re_plant_specs, storage_specs, dt_idx
+    ):
+        self.dt = (
+            pl.from_pandas(dt_idx.to_frame())
+            .lazy()
+            .select(pl.col("datetime").cast(pl.Datetime("us")))
+        )
+        on = pl.lit(1).alias("on")
+        dt = self.dt.with_columns(on)
+        combined_id = pl.concat_str(
+            [pl.col("plant_id_eia"), pl.col("generator_id")], separator="_"
+        ).alias("combined_id")
+        self.disp_convert = (
+            pl.from_pandas(dispatchable_specs.index.to_frame(), **self.schema)
+            .with_columns(combined_id)
+            .lazy()
+        )
+        self.disp_big_idx = (
+            self.disp_convert.with_columns(on).join(dt, on="on").drop("on").collect()
+        )
+        if re_plant_specs is not None:
+            self.re_convert = (
+                pl.from_pandas(re_plant_specs.index.to_frame(), **self.schema)
+                .with_columns(combined_id)
+                .lazy()
+            )
+            self.re_big_idx = (
+                self.re_convert.with_columns(on).join(dt, on="on").drop("on").collect()
+            )
+
+        self.storage_convert = (
+            pl.from_pandas(storage_specs.index.to_frame(), **self.schema)
+            .with_columns(combined_id)
+            .lazy()
+        )
+        self.storage_big_idx = (
+            self.storage_convert.with_columns(on).join(dt, on="on").drop("on").collect()
+        )
+
+    @property
+    def schema(self):
+        """Default schema overrides."""
+        return {"schema_overrides": {"plant_id_eia": pl.Int32}}
+
+    def from_pandas(self, df):
+        """Convert and apply schema to :class:`pandas.DataFrame`."""
+        return pl.from_pandas(df.reset_index(), **self.schema).fill_nan(None).lazy()
+
+    def __getstate__(self):
+        return None, {
+            name: getattr(self, name) for name in self.__slots__ if hasattr(self, name)
+        }
+
+    def __setstate__(self, state: tuple[Any, dict]):
+        _, state = state
+        for k, v in state.items():
+            if k in self.__slots__:
+                setattr(self, k, v)