Multiple functions in single udf

[dmpetrov]

Description

Claude Code like using this pattern: .map(a=fn1, b=fn2) instead of .map(a=fn1).map(b=fn2).

That's seem a good syntax sugar and needs to be supported.

• map

Updated:

- [ ] gen
- [ ] agg

• Optimization of a single map - map(a=f1, b=f2) should be squashed to a singe function without intermediate results map(c=f3) where f3 is like f2(f1()) and c is a superset of a and b
• Optimization of multiple maps - map(a=f1).map(b=f2) ==> map(c=f2(f1()))

[ilongin]

Adding fix for map is easy, but there are issues for gen and agg as for them multi-kwarg form doesn't have one obvious meaning

• gen: two generators produce different row counts. Concat (fn1's rows then fn2's rows)? Cartesian product? Sugar for chained .gen() (which changes fn2's input set)? None of these is clearly right.
• agg: same problem, plus partition_by - if fn1 and fn2 want different partition keys, that doesn't compose into one call.

I think we should keep gen/agg as-is (multi-kwarg stays an error) and fix this for map-only.

[dmpetrov]

we should keep gen/agg as-is

Yes, combining multiple gen and agg is a bit too advance technique. However, using map().get() in a single stage makes sense if we include a proper optimization - see below.

Optimization. It looks like I miss one part that's important here - these map(a=f1, b=f2) needs to be combined in a single call, not two separate calls. That's an optimization user expect - all lazy framework do this without saving intermediate result. We also need to apply this optimization for map(a=f1).map(b=f2) - to put them in a "single stage".

[ilongin]

we should keep gen/agg as-is

Yes, combining multiple gen and agg is a bit too advance technique. However, using map().get() in a single stage makes sense if we include a proper optimization - see below.

Optimization. It looks like I miss one part that's important here - these map(a=f1, b=f2) needs to be combined in a single call, not two separate calls. That's an optimization user expect - all lazy framework do this without saving intermediate result. We also need to apply this optimization for map(a=f1).map(b=f2) - to put them in a "single stage".

Ok, that is little bit different than my current opened PR which basically just did syntactic sugar map(a=f1, b=f2) -> .map(..).map(..) .. your suggestion does make much more sense, as it's real performance optimization...I will open new draft PR, I think what you suggeste should be easy to implement as well.