better errors for lazyframes (#1916)

Co-authored-by: Riccardo Cappuzzo <7548232+rcap107@users.noreply.github.com>

Author: jeromedockes

CHANGES.rst

@@ -38,6 +38,8 @@ Bug Fixes
   provided encoders is a scikit-learn Pipeline that starts with a skrub
   single-column transformer. :pr:`1899` by :user:`Jérôme Dockès <jeromedockes>`
   and :pr:`1900` by :user:`Jérôme Dockès <jeromedockes>`.
+- Errors raised when a polars LazyFrame is passed where an eager DataFrame is
+  expected are now clearer. :pr:`1916` by :user:`Jérôme Dockès <jeromedockes>`.
 
 Release 0.7.2
 =============

skrub/_check_input.py

@@ -64,20 +64,17 @@ def _check_not_pandas_sparse_pandas(df):
 
 
 def _check_is_dataframe(df):
+    if sbd.is_lazyframe(df):
+        warnings.warn(
+            "At the moment, skrub only works on eager DataFrames, calling collect()."
+        )
+        df = sbd.collect(df)
     if not sbd.is_dataframe(df):
         raise TypeError(
             "Only pandas and polars DataFrames are supported. Cannot handle X of"
             f" type: {type(df)}."
         )
-
-
-def _collect_lazyframe(df):
-    if not sbd.is_lazyframe(df):
-        return df
-    warnings.warn(
-        "At the moment, skrub only works on eager DataFrames, calling collect()."
-    )
-    return sbd.collect(df)
+    return df
 
 
 class CheckInputDataFrame(TransformerMixin, BaseEstimator):
@@ -122,7 +119,7 @@ def fit(self, X, y=None):
     def fit_transform(self, X, y=None):
         del y
         X = self._handle_array(X)
-        _check_is_dataframe(X)
+        X = _check_is_dataframe(X)
         self.module_name_ = sbd.dataframe_module_name(X)
         # TODO check schema (including dtypes) not just names.
         # Need to decide how strict we should be about types
@@ -133,13 +130,12 @@ def fit_transform(self, X, y=None):
         if sbd.column_names(X) != self.feature_names_out_:
             X = sbd.set_column_names(X, self.feature_names_out_)
         _check_not_pandas_sparse(X)
-        X = _collect_lazyframe(X)
         return X
 
     def transform(self, X):
         check_is_fitted(self, "module_name_")
         X = self._handle_array(X)
-        _check_is_dataframe(X)
+        X = _check_is_dataframe(X)
         module_name = sbd.dataframe_module_name(X)
         if module_name != self.module_name_:
             raise TypeError(
@@ -161,7 +157,6 @@ def transform(self, X):
         if sbd.column_names(X) != self.feature_names_out_:
             X = sbd.set_column_names(X, self.feature_names_out_)
         _check_not_pandas_sparse(X)
-        X = _collect_lazyframe(X)
         return X
 
     def _handle_array(self, X):

skrub/_dataframe/_common.py

@@ -502,7 +502,7 @@ def _shape_pandas(obj):
     return obj.shape
 
 
-@shape.specialize("polars")
+@shape.specialize("polars", argument_type=("DataFrame", "Column"))
 def _shape_polars(obj):
     return obj.shape
 

skrub/_dispatch.py

@@ -171,7 +171,7 @@ def _load_dataframe_module_info(name):
                 "module": polars,
                 "types": MappingProxyType(
                     {
-                        "DataFrame": (polars.DataFrame, polars.LazyFrame),
+                        "DataFrame": (polars.DataFrame,),
                         "LazyFrame": (polars.LazyFrame,),
                         "EagerFrame": (polars.DataFrame,),
                         "Column": (polars.Series,),
@@ -262,6 +262,12 @@ def raise_dispatch_unregistered_type(obj, kind="object"):
 evaluated) rather than computing it immediately.
  """
         )
+    if type(obj).__name__ == "LazyFrame":
+        raise TypeError(
+            "Expected a DataFrame but got an object "
+            f"of type {type(obj)}. LazyFrames are not yet supported; "
+            "please call collect() to convert to a DataFrame."
+        )
     raise TypeError(
         "Operation not supported on this object. Expecting a Pandas or Polars "
         f"{kind}, but got an object of type {type(obj)}."

skrub/_reporting/_table_report.py

@@ -242,11 +242,6 @@ def __init__(
         self.dataframe = (
             sbd.to_frame(dataframe) if sbd.is_column(dataframe) else dataframe
         )
-        if sbd.is_polars(dataframe) and sbd.is_lazyframe(dataframe):
-            raise ValueError(
-                "The TableReport does not support lazy dataframes. Please call"
-                " `.collect()` to use the TableReport on the current dataframe."
-            )
         self.n_columns = sbd.shape(self.dataframe)[1]
 
     def _set_minimal_mode(self):

skrub/_reporting/tests/test_table_report.py

@@ -103,9 +103,7 @@ def test_lazyframe_exception():
     pl = pytest.importorskip("polars")
     lazy_df = pl.DataFrame({"a": ["1", "2", "3"]}).lazy()
 
-    with pytest.raises(
-        ValueError, match=r"TableReport does not support lazy dataframes"
-    ):
+    with pytest.raises(TypeError, match=r".*LazyFrames are not yet supported"):
         TableReport(lazy_df)