diff --git a/python/hsfs/engine/python.py b/python/hsfs/engine/python.py
index 510bf5223f..806725d089 100644
--- a/python/hsfs/engine/python.py
+++ b/python/hsfs/engine/python.py
@@ -930,7 +930,12 @@ def split_labels(
     ]:
         if labels:
             labels_df = df[labels]
-            df_new = df.drop(columns=labels)
+            if HAS_POLARS and (
+                    isinstance(df, pl.DataFrame) or isinstance(df, pl.dataframe.frame.DataFrame)
+            ):
+                df_new = df.drop(labels)
+            else:
+                df_new = df.drop(columns=labels)
             return (
                 self._return_dataframe_type(df_new, dataframe_type),
                 self._return_dataframe_type(labels_df, dataframe_type),
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 058b967105..778127436d 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -94,7 +94,7 @@ dev-pandas1 = [
 ]
 dev = ["hopsworks[dev-no-opt,great-expectations,polars]"]
 polars=[
-    "polars>=0.20.18,<=0.21.0",
+    "polars>=1.8.2,<=1.31.0",
     "pyarrow>=17.0",
 ]