Add DataFrame::map utility .map function for DataFrame for modifying internal LogicalPlan

[phisn]

Is your feature request related to a problem or challenge?

Currently when applying custom LogicalPlans (extensions or transformation functions) on DataFrames I need to convert DataFrame to a LogicalPlan and then back to a DataFrame making cumbersome code.

Describe the solution you'd like

It would be very helpful to have a .map function which transforms the LogicalPlan inside DataFrames without reconstructing it.

Describe alternatives you've considered

No response

Additional context

No response

[alamb]

This seems like a good idea to me

As I understand it, it would allow things like

let df = ctx.sql("SELECT * from foo");
let df = df.map(my_awesome_rewrite)?;
...

fn my_awesome_plan(plan: LogicalPlan| ) -> Result<LogicalPlan> {
   ...
  Ok(new_plan)
}

Any thoughts @timsaucer or @Omega359 ?

[timsaucer]

We have something similar in datafusion-python here. It lets you do something like what you're describing but it operates on a dataframe and not a logical plan.

In python you can currently do something like this (this is from our unit tests)

    def add_string_col(df_internal) -> DataFrame:
        return df_internal.with_column("string_col", literal("string data"))

    def add_with_parameter(df_internal, value: Any) -> DataFrame:
        return df_internal.with_column("new_col", literal(value))

    df = df.transform(add_string_col).transform(add_with_parameter, 3)

If you had a DataFrame df then this would add in two more columns. It's a trivial example, but it allows for some really nice chaining of operations like in the last line. I would definitely support adding something like that on the rust side.

But what I'm doing here doesn't exactly match up with what the issue requests, which is to work on the LogicalPlan. So we could do something similar but df.transform_plan that will do the conversion from dataframe to plan, execute the function/closure, and convert back to dataframe.

@phisn would that meet your needs?

[Omega359]

I personally haven't had the need to go into a LogicalPlan from a dataframe and back again but I could see it being useful.

[phisn]

@timsaucer My specific use case is what @alamb described. The problem with the transform approach is that I am forced to .logical_plan().clone as well as create a new DataFrame each time, since it's about directly manipulating the LogicalPlan.

The problem arises with applying extensions and using functions that are not always called using a DataFrame. For example when the context is not always available.