support polars

pycaret/internal/preprocess/preprocessor.py

@@ -1,10 +1,13 @@
+# Copyright (C) 2019-2024 PyCaret
 # Author: Mavs (m.524687@gmail.com)
+# Contributors (https://github.com/pycaret/pycaret/graphs/contributors)
 # License: MIT
 
 from copy import deepcopy
 
 import numpy as np
 import pandas as pd
+import polars as pl
 from category_encoders.basen import BaseNEncoder
 from category_encoders.one_hot import OneHotEncoder
 from category_encoders.ordinal import OrdinalEncoder
@@ -118,6 +121,10 @@ def _prepare_dataset(self, X, y=None):
         """
         self.logger.info("Set up data.")
 
+        # Convert Polars DataFrame to Pandas DataFrame if it's Polars
+        if isinstance(X, pl.DataFrame):
+            X = X.to_pandas()
+
         # Make copy to not overwrite mutable arguments
         X = to_df(deepcopy(X))
 

requirements-optional.txt

@@ -40,3 +40,6 @@ fugue
 flask
 Werkzeug>=2.2,<3.0
 triad<0.9.2 # fix pandas enforce_type error (https://github.com/fugue-project/fugue/issues/526)
+
+#Dataframe
+polars

tests/test_datasets.py

@@ -1,7 +1,9 @@
 import pandas as pd
+import polars as pl
 import pytest
 
 from pycaret.datasets import get_data
+from pycaret.internal.preprocess import Preprocessor
 
 
 def test_datasets():
@@ -36,6 +38,28 @@ def test_datasets():
     assert cols == 24
     assert data.size == 576000
 
+    ################################
+    # Conversion from Polars to Pandas ####
+    ################################
+
+    # creating a Polars DataFrame manually
+    df = pl.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5],
+            "fruits": ["banana", "banana", "apple", "apple", "banana"],
+            "B": [5, 4, 3, 2, 1],
+            "cars": ["beetle", "audi", "beetle", "beetle", "beetle"],
+        }
+    )
+
+    # testing conversion
+    preprocessor = Preprocessor()
+    data_converted = preprocessor._prepare_dataset(df)
+    assert isinstance(data_converted, pd.DataFrame)
+    rows, cols = data_converted.shape
+    assert rows == 5
+    assert cols == 4
+
     ################################
     # GitHub Specific folder ####
     ################################