Support for Polars and other dataframes

chainladder/core/base.py

@@ -2,6 +2,8 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
 import pandas as pd
+from packaging import version
+
 import numpy as np
 from chainladder.utils.cupy import cp
 from chainladder.utils.sparse import sp
@@ -402,7 +404,30 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             return obj
         else:
             raise NotImplementedError()
-
+
+    def _interchange_dataframe(self, data):
+        """
+        Convert an object supporting the __dataframe__ protocol to a pandas DataFrame.
+        Requires pandas version > 1.5.2.
+
+        Parameters
+        ----------
+        data :
+            Dataframe object supporting the __dataframe__ protocol.
+        Returns
+        -------
+        out: pd.DataFrame
+            A pandas DataFrame.
+        """
+        # Check if pandas version is greater than 1.5.2
+        if version.parse(pd.__version__) >= version.parse("1.5.2"):
+            return pd.api.interchange.from_dataframe(data)
+
+        else:
+            # Raise an error prompting the user to upgrade pandas
+            raise NotImplementedError("Your version of pandas does not support the DataFrame interchange API. "
+                                    "Please upgrade pandas to a version greater than 1.5.2 to use this feature.")
+
     def __array_function__(self, func, types, args, kwargs):
         from chainladder.utils.utility_functions import concat
 

chainladder/core/tests/test_triangle.py

@@ -1,5 +1,6 @@
 import chainladder as cl
 import pandas as pd
+import polars as pl
 import numpy as np
 import copy
 import pytest
@@ -745,6 +746,9 @@ def test_halfyear_development():
         ["2012-01-01", "2013-12-31", "incurred", 200.0],
     ]
 
+    df_polars = pl.DataFrame(data)
+    df_polars.columns = ["origin", "val_date", "idx", "value"]
+
     assert (
         type(
             cl.Triangle(
@@ -757,3 +761,35 @@ def test_halfyear_development():
             )
         )
     ) == cl.Triangle
+
+    assert (
+        type(
+            cl.Triangle(
+                data=df_polars,
+                index="idx",
+                columns="value",
+                origin="origin",
+                development="val_date",
+                cumulative=True,
+            )
+        )
+    ) == cl.Triangle
+
+    assert (
+            cl.Triangle(
+                data=pd.DataFrame(data, columns=["origin", "val_date", "idx", "value"]),
+                index="idx",
+                columns="value",
+                origin="origin",
+                development="val_date",
+                cumulative=True,
+            ) ==
+            cl.Triangle(
+                data=df_polars,
+                index="idx",
+                columns="value",
+                origin="origin",
+                development="val_date",
+                cumulative=True,
+            )
+    )

chainladder/core/triangle.py

@@ -123,7 +123,9 @@ def __init__(
     ):
         if data is None:
             return
-
+        elif not isinstance(data, pd.DataFrame) and hasattr(data, "__dataframe__"):
+            data = self._interchange_dataframe(data)
+
         index, columns, origin, development = self._input_validation(
             data, index, columns, origin, development
         )
@@ -270,7 +272,7 @@ def __init__(
             self.ddims = obj.ddims
             self.values = obj.values
             self.valuation_date = pd.Timestamp(options.ULT_VAL)
-
+    
     @staticmethod
     def _split_ult(data, index, columns, origin, development):
         """Deal with triangles with ultimate values"""

ci/environment-latest.yaml

@@ -6,6 +6,7 @@ channels:
 
 dependencies:
   - conda-forge::pandas>=2.0
+  - polars
   - scikit-learn>=0.23
   - sparse>=0.9
   - numba

ci/environment-test.yaml

@@ -7,6 +7,7 @@ channels:
 
 dependencies:
   - pandas>=0.23
+  - polars>=0.16.2
   - scikit-learn>=0.23
   - sparse>=0.9
   - numba

environment-dev.yaml

@@ -14,6 +14,7 @@ dependencies:
   - ipykernel
 
   - pandas
+  - polars
   - scikit-learn
   - sparse
   - numba

environment-docs.yaml

@@ -11,6 +11,7 @@ dependencies:
   - git
 
   - pandas
+  - polars
   - scikit-learn
   - sparse
   - numba