Refactor POC analysis module to use POLARS instead of pandas. FOR MAC…

…OS: you’ll need to install polars via pip install polars-lts-cpu.

may have to delete,
pip uninstall polars
pip install polars-lts-cpu
(we also use pandera[ioc,polars]

openenergyid/pandera_poc/analysis.py

@@ -1,11 +1,11 @@
 """Proof of concept of a data analysis module using pandera."""
 
-import pandas as pd
+import polars as pl
 
 
 def analyse(
-    df: pd.DataFrame,
-) -> pd.DataFrame:
+    lf: pl.LazyFrame,
+) -> pl.LazyFrame:
     """Perform analysis on the input data and return the output data."""
     # Perform analysis
-    return df
+    return lf

openenergyid/pandera_poc/models.py

@@ -1,38 +1,38 @@
 """Example of a Pandera schema for input and output data validation."""
 
-import pandera as pa
+import pandera.polars as pap
+import polars as pl
+# from pandera.typing import Series
 
-from pandera.typing import Series, Index
 
-
-class InputModel(pa.DataFrameModel):
+class InputModel(pap.DataFrameModel):
     """Pandera schema for input data validation."""
 
-    index: Index[int] = pa.Field(ge=0)
-    column1: Series[int] = pa.Field(le=10)
-    column2: Series[float] = pa.Field(lt=-1.2)
-    column3: Series[str] = pa.Field(str_startswith="value_")
+    # index: Optional[Index[int]]
+    column1: int = pap.Field(le=10)
+    column2: pl.Float64 = pap.Field(lt=-1.2)
+    column3: str = pap.Field(str_startswith="value_")
 
-    @pa.check("column3")
+    @pap.check("column3")
     @classmethod
-    def column_3_check(cls, series: Series[str]) -> bool:
+    def column_3_check(cls, series: pl.Series) -> bool:
         """Check that column3 values have two elements after being split with '_'"""
-        return series.str.split("_", expand=True).shape[1] == 2
+        return len(series.str.split("_")) == 2
 
 
-class OutputModel(pa.DataFrameModel):
+class OutputModel(pap.DataFrameModel):
     """Pandera schema for output data validation."""
 
-    index: Index[int] = pa.Field(ge=0)
-    column1: Series[int] = pa.Field(le=10)
-    column2: Series[float] = pa.Field(lt=-1.2)
-    column3: Series[str] = pa.Field(str_startswith="value_")
+    # index: pl.int = pap.Field(ge=0)
+    column1: int = pap.Field(le=10)
+    column2: float = pap.Field(lt=-1.2)
+    column3: str = pap.Field(str_startswith="value_")
 
-    @pa.check("column3")
+    @pap.check("column3")
     @classmethod
-    def column_3_check(cls, series: Series[str]) -> bool:
+    def column_3_check(cls, series: pl.Series) -> bool:
         """Check that column3 values have two elements after being split with '_'"""
-        return series.str.split("_", expand=True).shape[1] == 2
+        return len(series.str.split("_")) == 2
 
     class Config:
         """Pandera schema configuration."""