feat(rust,python): Add GPU support to the LazyFrame profiler

[Matt711]

Closes #20039

[codecov]

Codecov Report

Attention: Patch coverage is 36.84211% with 60 lines in your changes missing coverage. Please review.

Project coverage is 79.98%. Comparing base (45ec22a) to head (0be4de8).

Files with missing lines	Patch %	Lines
crates/polars-python/src/lazyframe/general.rs	33.33%	22 Missing ⚠️
...olars-mem-engine/src/executors/scan/python_scan.rs	46.87%	17 Missing ⚠️
crates/polars-lazy/src/frame/mod.rs	0.00%	10 Missing ⚠️
py-polars/polars/lazyframe/frame.py	42.85%	6 Missing and 2 partials ⚠️
crates/polars-expr/src/state/node_timer.rs	50.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #20693      +/-   ##
==========================================
- Coverage   79.99%   79.98%   -0.02%     
==========================================
  Files        1598     1598              
  Lines      229199   229285      +86     
  Branches     2620     2623       +3     
==========================================
+ Hits       183352   183387      +35     
- Misses      45248    45297      +49     
- Partials      599      601       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ritchie46]

We should use the enter_polars which handles the allow_threads.

[ritchie46]

Nit: Can we make this query multiline? Something like:

(
    lf.group_by("a", maintain_order=True)
    .agg(
         pl.all().sum()
      ).sort("a")
    .profile
)

[ritchie46]

I believe this code is exactly the same as in collect. Can we put it in a function?

[ritchie46]

I don't think it is needed to leak those internals. See comment: crates/polars-mem-engine/src/executors/scan/python_scan.rs

[ritchie46]

Instead of leaking our timer nodes which, I really don't like as now we cannot refactor as CuDF is dependent on internals other than our DSL.

Instead the python callable could accept an argument profile: bool and if set return a list of timing tuples:

[(operation: str,, start: time, end: time)]

We can then unpack those tuples here and update the NodeTime accordingly.

[Matt711]

Thanks for the review @ritchie46. There might be a casting problem with this approach. Is there a way to convert a u64 to a std::time::Instant?

[ritchie46]

Instead of storing the Instants we could convert them to u64 immediately in the store method.

[Matt711]

Once I time the node in cuDF, I'll get two u64 values. The store method of NodeTimer takes Instants. I dont think there's a way to convert a u64 to Instant. Is that right?

[ritchie46]

We can bypass the store method for GPU. And then store u64s instead of Instants.

[Matt711]

I'm giving your suggestion a try @ritchie46

1. How would I unpack the return call? I have some thing like

                let args = (
                    python_scan_function,
                    with_columns.map(|x| x.into_iter().map(|x| x.to_string()).collect::<Vec<_>>()),
                    predicate,
                    n_rows,
                    true,
                );
                callable.call1(args).map_err(to_compute_err) <--- returns tuple[DataFrame, tuple[str, int, int]] in python

2. How would I bypass the store method which relies upon std::time::Instants?

[ritchie46]

You can use pyo3 to access the items of a tuple and then extract the proper types. (See later in this file how we can get the DataFrame).

How would I bypass the store method which relies upon std::time::Instants?

Make a store_raw method.

[Matt711]

Update: I'm able to extract the items from the PyTuple returned from the python callback. I also added the raw_store to NodeTimer. Here's what profiling with the GPU looks like so far:

In [1]: import polars as pl

In [2]: lf = pl.LazyFrame(
   ...: 
   ...:     {
   ...: 
   ...:         "a": ["a", "b", "a", "b", "b", "c"],
   ...: 
   ...:         "b": [1, 2, 3, 4, 5, 6],
   ...: 
   ...:         "c": [6, 5, 4, 3, 2, 1],
   ...: 
   ...:     }
   ...: 
   ...: )
   ...: 
   ...: lf.group_by("a", maintain_order=True).agg(pl.all().sum()).sort(
   ...: 
   ...:     "a"
   ...: 
   ...: ).profile(engine="gpu")
Out[2]: 
(shape: (3, 3)
 ┌─────┬─────┬─────┐
 │ a   ┆ b   ┆ c   │
 │ --- ┆ --- ┆ --- │
 │ str ┆ i64 ┆ i64 │
 ╞═════╪═════╪═════╡
 │ a   ┆ 4   ┆ 10  │
 │ b   ┆ 11  ┆ 10  │
 │ c   ┆ 6   ┆ 1   │
 └─────┴─────┴─────┘,
 shape: (2, 3)
 ┌──────────────┬──────────────────┬──────────────────┐
 │ node         ┆ start            ┆ end              │
 │ ---          ┆ ---              ┆ ---              │
 │ str          ┆ u64              ┆ u64              │
 ╞══════════════╪══════════════════╪══════════════════╡
 │ optimization ┆ 0                ┆ 9591759057601683 │
 │ sort         ┆ 9591759057601683 ┆ 9591759094456985 │
 └──────────────┴──────────────────┴──────────────────┘)
***The times are in nanoseconds not microseconds

There are a couple of problems with the result that I'm encountering.

1. The NodeTimer keeps a query_start attribute (a std::time::Instant) which I cannot use to adjust all of the raw u64 times. Do you have any ideas on capturing the start time and having it available in NodeTimer::finalize?
2. If you look at the output above, only the final sort operation is being captured. There's probably a copy happening somewhere so that I'm not updating the same NodeTimer object each time I call store_raw. I'm surprised because the execution state is mutable. Do I need some sort of global NodeTimer to resolve this?

[wence-]

Tried an alternative approach in #21534.

[Matt711]

Closing in favor of #21534. Thanks for your help on this @ritchie46!