feat: "carefully" allow for dask Expr that modify index

[FBruzzesi]

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

If you have comments or can explain your changes, please do so below.

Pretty dangerous stuff to workaround the dask index.

To assess that the implementation is working as expected, I implemented both sort (different index but same length) and drop_nulls (different index due to different length)

[FBruzzesi]

Sadly this broadcast is not working as drop_nulls does not return a scalar. I would consider this an edge case and focus on the broader support

[anopsy]

[MarcoGorelli]

thanks for trying this - i'll test it out and see if there's a perf impact

[FBruzzesi]

this is guaranteed to not end up in StopIteration error as if everything was a scalar the previous block would have been entered and returned

[MarcoGorelli]

we've got the notebooks in tpch/notebooks, the first two support Dask - fancy running them with this branch and seeing if there's any perf impact?

[FBruzzesi]

Hey @MarcoGorelli, I am giving another thought on this feature (which I would still love to see), here is a simple idea to have partial support without loss of performance:

• sort is the only method that changes the index and result in an output with the same length. Instead of changing the index to each series, we can do that specifically in sort, namely by assigning the original index.

  def sort(self: Self, *, descending: bool = False, nulls_last: bool = False) -> Self:
       na_position = "last" if nulls_last else "first"
  
       def func(_input: Any, ascending: bool, na_position: bool) -> Any:  # noqa: FBT001
           name = _input.name
  
           result =_input.to_frame(name=name).sort_values(
               by=name, ascending=ascending, na_position=na_position
           )[name]
           return de._expr.AssignIndex(result, _input.index)
  
       return self._from_call(
           func,
           "sort",
           not descending,
           na_position,
           returns_scalar=False,
       )

• All the other methods that change the index, do so by reducing the length of the series. In my working experience and in TPCH queries they are mostly used before a reduction or in isolation, therefore we should not worry of changing their index. Example:

  df.select(
      head_sum=pl.col("a").head().sum(),
      tail_mean=pl.col("a").tail().mean(),
  )

• What is left and unsupported you may ask? Multiple reductions operations ending up with the same length, different from the original, won't be possible. Example:

  df.select(
      head=pl.col("a").head(),
      tail=pl.col("a").tail(),
  )

What do you think?

[FBruzzesi]

@MarcoGorelli I am tagging this as ready for review as I re-worked it a bit more.

The TL;DR is:

• sort is kind of special, as it modifies the index but returns a Series of the same length of the original one, therefore in such specific case I am manually re-assigning the index
• for all other methods, I added a boolean flag to DaskExpr called modifies_index and:
  • that is not allowed in with_columns
  • in select it should be allowed only if there are no other exprs or there is a reduction following (I need to address both these cases actually).

Yet before developing further, I would like some feedback on how likable this approach is and if we want to move forward with it 🙏🏼

[FBruzzesi]

So... head has a npartitions param which can be set to -1 and scan all partitions, while tail does not. This means that if we have multiple partitions, then this implementation of tail may not be what we expect

[phofl]

I haven't seen cases where users actually want more than 1 partition if they call head or tail tbh, Yes this is technically an issue, but not something I've encountered in the wild

[MarcoGorelli]

thanks @FBruzzesi !

to be honest I don't know about using such private methods, it makes me feel slightly uneasy - @phofl do you have time/interest in taking a look? specifically the de._collection.Series(de._expr.AssignIndex(result, _input.index)) part in narwhals/_dask/expr.py

I think that for sql engines (like duckdb, which hopefully we can get to eventually) operations like df.select(nw.col('a').sort(), nw.col('b')) would be problematic anyway, so I don't think it'd be an issue to leave them out of the Narwhals area of support

[phofl]

Yeah I can second @MarcoGorelli here, please don't do this. We are sorting a Series (i.e. a single column from the df), correct?

I would:

tmp = _input.reset_index().sort_values(...)
result = tmp[_input.name]
result.index = tmp["use the index name of _input"]

[phofl]

Or do you want to keep the Index of the input?

If yes, this is fundamentally a bad idea in Dask, it will shoot you in the foot all over the place. You have zero guarantees that the partitions keep their lengths when sorting (it is a lot more likely that they do not), so this is bound to fail in all kinds of places

[FBruzzesi]

Hey @phofl thanks for taking the time.

We are sorting a Series (i.e. a single column from the df), correct?

Yes indeed, but with the final goal to potentially add it as a new column to the original dataframe, and that's where the index misalignment gets into the game.

I will try the approach you are suggesting, which is not too far off from what is already implemented, and see if everything else falls into place 🙏🏼

Edit: it just ends up raising:

AssertionError: value needs to be aligned with the index

Traceback

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[60], line 3
      1 tmp = df_dd["a"].to_frame(name="a").sort_values("a")
      2 result = tmp["a"]
----> 3 result.index = df_dd["a"].index

    655 @index.setter
    656 def index(self, value):
--> 657     assert expr.are_co_aligned(
    658         self.expr, value.expr
    659     ), "value needs to be aligned with the index"
    660     _expr = expr.AssignIndex(self, value)
    661     self._expr = _expr

AssertionError: value needs to be aligned with the index

[FBruzzesi]

Hi @phofl, apologies to call you in the mix once more.

I have a few questions in order to make this work and guarantee that we don't end up with a

fundamentally a bad idea in Dask, it will shoot you in the foot all over the place.

• How can we test for when Dask will shoot us in the foot if we do something bad?
• The latest approach TL;DR is that if a method changes the index, then it either has to be followed by a reduction or be a single selection. Examples:
  • Reductions:

    df.select(
        head_sum=pl.col("a").head().sum(),
        tail_mean=pl.col("a").tail().mean(),
    )

    which would translate to something like dd.concat([df["a"].head().sum(), df["a"].tail().mean()])
  • Single selection:

     df.select(
        head=pl.col("a").head(),
     )

  In sight of the first question, what do you think about this approach? Is it a fundamentally bad idea?