Implement pl.DataFrame.collect similarly to pl.LazyFrame.lazy

[StefanBRas]

Problem description

In the same way pl.LazyFrame has .lazy() it would be nice if pl.DataFrame had .collect(), so it's easier to write code that works for both. (as the docstring for pl.LazyFrame.lazy notes).

My specific motivation is written some generic test assertions where I now have to do:

def assert_no_new_rows(new: pl.DataFrame | pl.LazyFrame, old: pl.DataFrame | pl.LazyFrame):
    if isinstance(new, pl.LazyFrame):
        new = new.collect()
    if isinstance(old, pl.LazyFrame):
        old = old.collect()
    assert new.height == old.height

If pl.DataFrame.collect was Implemented this could be written as:

def assert_no_new_rows(new: pl.DataFrame | pl.LazyFrame, old: pl.DataFrame | pl.LazyFrame):
   assert new.collect().height == old.collect().height

If you think it makes sense, I can make a PR.

[s-banach]

Maybe your use of DataFrame | LazyFrame makes more sense outside of this toy example.
But this example is nonsensical, because your function is calling .collect() and then discarding the result, meaning you'll have to call .collect() again outside the function, and run the whole computation twice.

[StefanBRas]

I wouldn't call it nonsensical, it's the exact functionality I want. It's for testing correctness and it's fast enough that I'd prefer to (maybe) have to do a recalculation as opposed to make conditional paths for DataFrame versus LazyFrame.

I feel like the argument is the exact same as the argument for LazyFrame.lazy - "Useful for writing code that expects either a DataFrame or LazyFrame."

Another example could be that I wanted to plot some intermediate subresults:

def plot(df: pl.LazyFrame | pl.DataFrame):
    data = df.select('value').filter(pl.col('value') > 100).collect()
    my_plot(data)

def calculate_and_plot_thing(df: pl.LazyFrame | pl.DataFrame):
    first_df = df.with_columns(...)
    plot(first_df)
    second_df = first_df.with_columns(...)
    plot(second_df)
    return second_df.with_columns(...)

If I rewrote plot to only take DataFrame I would have to collect outside plot meaning that I would either have to collect without the optimization from select and filter or I would have to put the select and filter outside of the plot function, making it less obvious and modular.

[s-banach]

Sorry for being rude. I didn't even know about LazyFrame.lazy(), which shows there is some precedent for this kind of thing.

[ghuls]

df.lazy().collect() should work for both DataFrame and LazyFrame.

[StefanBRas]

df.lazy().collect() should work for both DataFrame and LazyFrame.

Ah yes, that works. Thanks!

Given the existence of a lazy: (DataFrame | LazyFrame) -> LazyFrame I think it makes sense to have the inverse collect: (LazyFrame | DataFrame) -> DataFrame but I guess it's matter of preference. There are the caveats of potential extra work if people end up collecting the same LazyFrame multiple times.

[stinodego]

Closing in favor of #7882