I haven't been particularly good about this myself, but it's my impression that lambdas and closures within functions are more expensive to serialize than top-level functions. I wish I could find some actual documentation about this, but the closest I've seen is this comment from Jim: #6192 (comment)

I've heard that too, but don't really know the reasoning. Certainly serializing an arbitrary function is more expensive than something defined in a module, because the one in a module can just be referenced by name, instead of inlining all the code. @jcrist?

Here though, if you're taking partition_info, orig_funcion is very, very likely something a user wrote, and probably isn't from library code, so probably needs its code serialized anyway. So in this specific case, I wonder if it's less likely to be a concern?

When I wrote this, I did originally try to avoid the lambda, and went in all sorts of blockwise circles trying to figure it out. But Blockwise only supports positional arguments (kwargs are just literals), and trying to change that was not easy. And partition_info may have to be given to the function as a kwarg. So in the end, something like this seemed the easiest.

I guess we could instead do:

# module-level
def apply_add_partition_info(func, partition_info, *args, **kwargs):
    return func(*args, **kwargs, partition_info=partition_info)

...
if has_keyword(func, "partition_info"):
    func = functools.partial(apply_add_partition_info, func)

This is an interesting application of BlockwiseDep. I don't think it's quite what was envisioned for the helper class, but it doesn't seem obviously wrong, and the win in this case is pretty clear. I wonder if it would be worth renaming that class or improving the docs, since it is not in this case strictly used for dependencies?

As far as I can tell, this is very similar to what @bmerry was trying in #7686 (though that one seems lazier, which may have some benefits). Unfortunately, that PR seems to have not gotten much attention, though it should have,

Yup, the lack of laziness here was just laziness on my part :). This could be its own BlockwiseDep subclass that generates the info at construction time. Though since divisions would have to be embedded anyway, I felt like the savings from lazily generating i just weren't worth the effort.

I agree that this is a perhaps-unintended but very useful application of BlockwiseDep. I have a feeling that this is a pattern we'll want to reuse in many places. This is why I felt that #7513 would be powerful: it gives you an easy interface for passing in auxiliary data that needs to be broadcast like everything else, but is literal instead of delayed. I have a feeling that with a few basic constructs and the ability to pass in arrays/lists/iterables wrapped in a marker class, you could compose pretty much anything you need without defining new BlockwiseDep subclasses.

For example, I'd want to use this for #8445, in order to pass in the filenames here:

And yes, this seems pretty much exactly like what @bmerry was trying! Including the fact that you added support for passing a BlockwiseDep instance into array.blockwise.blockwise, which I also did in dataframe.core.partitionwise_graph, which to me is an equivalent function.

I wonder if it would be worth renaming that class or improving the docs, since it is not in this case strictly used for dependencies?

Maybe BlockwiseLiteral? BlockwiseArg? I also perhaps like the idea of users not having to interact with it at all in basic cases—what if map_partitions/map_blocks automatically supported sequences or array-likes? (Just judging this on whether shapes align might be too much though; I could see wanting some explicit wrapping to clarify whether it's a literal or a broadcastable.) That plus block_info might get you everything you'd need as a user.

I've heard that too, but don't really know the reasoning. Certainly serializing an arbitrary function is more expensive than something defined in a module, because the one in a module can just be referenced by name, instead of inlining all the code. @jcrist?

Here though, if you're taking partition_info, orig_funcion is very, very likely something a user wrote, and probably isn't from library code, so probably needs its code serialized anyway. So in this specific case, I wonder if it's less likely to be a concern?

When I wrote this, I did originally try to avoid the lambda, and went in all sorts of blockwise circles trying to figure it out. But Blockwise only supports positional arguments (kwargs are just literals), and trying to change that was not easy. And partition_info may have to be given to the function as a kwarg. So in the end, something like this seemed the easiest.

I guess we could instead do:

# module-level
def apply_add_partition_info(func, partition_info, *args, **kwargs):
    return func(*args, **kwargs, partition_info=partition_info)

...
if has_keyword(func, "partition_info"):
    func = functools.partial(apply_add_partition_info, func)

Yup, the lack of laziness here was just laziness on my part :). This could be its own BlockwiseDep subclass that generates the info at construction time. Though since divisions would have to be embedded anyway, I felt like the savings from lazily generating i just weren't worth the effort.

I agree that this is a perhaps-unintended but very useful application of BlockwiseDep. I have a feeling that this is a pattern we'll want to reuse in many places. This is why I felt that #7513 would be powerful: it gives you an easy interface for passing in auxiliary data that needs to be broadcast like everything else, but is literal instead of delayed. I have a feeling that with a few basic constructs and the ability to pass in arrays/lists/iterables wrapped in a marker class, you could compose pretty much anything you need without defining new BlockwiseDep subclasses.

For example, I'd want to use this for #8445, in order to pass in the filenames here:

And yes, this seems pretty much exactly like what @bmerry was trying! Including the fact that you added support for passing a BlockwiseDep instance into array.blockwise.blockwise, which I also did in dataframe.core.partitionwise_graph, which to me is an equivalent function.

I wonder if it would be worth renaming that class or improving the docs, since it is not in this case strictly used for dependencies?

Maybe BlockwiseLiteral? BlockwiseArg? I also perhaps like the idea of users not having to interact with it at all in basic cases—what if map_partitions/map_blocks automatically supported sequences or array-likes? (Just judging this on whether shapes align might be too much though; I could see wanting some explicit wrapping to clarify whether it's a literal or a broadcastable.) That plus block_info might get you everything you'd need as a user.