refactored lazy frames to use .collect_schema().names() to fix warning

sktime · Aug 15, 2024 · e8cacd6 · e8cacd6
1 parent c220f37
commit e8cacd6
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/skpro/datatypes/_adapter/polars.py b/skpro/datatypes/_adapter/polars.py
@@ -34,7 +34,10 @@ def check_polars_frame(obj, return_metadata=False, var_name="obj", lazy=False):
     if _req("n_features", return_metadata):
         metadata["n_features"] = obj.width
     if _req("feature_names", return_metadata):
-        metadata["feature_names"] = obj.columns
+        if lazy:
+            metadata["feature_names"] = obj.collect_schema().names()
+        else:
+            metadata["feature_names"] = obj.columns
 
     # check if there are any nans
     #   compute only if needed