Partition_by should returns tuples as keys when partitioning on a list of a single column, like group_by does

[CaselIT]

Checks

• [X] I have checked that this issue has not already been reported.

• [X] I have confirmed this bug exists on the latest version of Polars.

Reproducible example

import polars as pl

df = pl.DataFrame({"a": [1, 1, 2, 2], "b": range(4)})

keys = [k for k, _ in df.group_by("a", maintain_order=True)]
print(keys)
keys = [k for k, _ in df.group_by(["a"], maintain_order=True)]
print(keys)
keys = [k for k, _ in df.group_by(["a", "b"], maintain_order=True)]
print(keys)
keys = list(df.partition_by("a", maintain_order=True, as_dict=True).keys())
print(keys)
keys = list(df.partition_by(["a"], maintain_order=True, as_dict=True).keys())
print(keys)
keys = list(df.partition_by(["a", "b"], maintain_order=True, as_dict=True).keys())
print(keys)

Log output

[1, 2]
[(1,), (2,)]
[(1, 0), (1, 1), (2, 2), (2, 3)]
[1, 2]
[1, 2]
[(1, 0), (1, 1), (2, 2), (2, 3)]

Issue description

Passing a list of a single column to partition by should return keys as one element tuple, similarly to what iterating on a groupby does.

The current behaviour is inconsistent with how partition by behaves with more than one by column passed, making hard to generalize code that uses is if this takes the partition columns from the outside.

Expected behavior

When passed a list of columns to partition_by should always return keys as tuple, like how list(group_by) does

Edit

This bahaviours seems to happen also when using selectors in partition by, in particular depending on the number of column matched at runtime by the selector a scalar or a select is returned. This is not great too, but may be another issue. Mentioned here: https://github.com/pola-rs/polars/issues/13371#issuecomment-1880684263

Installed versions

--------Version info---------
Polars:               0.20.2
Index type:           UInt32
Platform:             Windows-10-10.0.19045-SP0
Python:               3.10.13 | packaged by Anaconda, Inc. | (main, Sep 11 2023, 13:24:38) [MSC v.1916 64 bit (AMD64)]

----Optional dependencies----
adbc_driver_manager:  <not installed>
cloudpickle:          2.2.1
connectorx:           <not installed>
deltalake:            <not installed>
fsspec:               <not installed>
gevent:               <not installed>
matplotlib:           3.8.0
numpy:                1.26.2
openpyxl:             3.1.2
pandas:               2.1.3
pyarrow:              14.0.1
pydantic:             1.10.13
pyiceberg:            <not installed>
pyxlsb:               <not installed>
sqlalchemy:           <not installed>
xlsx2csv:             <not installed>
xlsxwriter:           3.1.9

[deanm0000]

I think this should just be a matter of checking if by is a list or tuple then in the len(by)==1 case return (p[by[0]][0],)

[CaselIT]

Note that maybe it's intended to be like this, but at that point group_by should match the behaviour. The main issue is the inconsistency between the two

[gab23r]

I agree, I already raise this behavior here: #11569

[CaselIT]

I actually skimmed that one, but though it was about something else, sorry.

In my personal opinion the better option here is to return a tuple in this case, but if its more consistent to return a scalar by considering by='a' and by=['a'] as equivalent then I believe that group_by should too. (it would break compatibility with pandas, that behaves like the current group by does, treating by='a' and by=['a'] as different)

edit: pandas df.groupby().groups that returns a similarly keyed dict behaves as partition_by, having scalars for both by='a' and by=['a']

[gab23r]

IMHO, by='a' and by=['a'] are not equivalent, and I would prefer the output types to be align to the input types and not on the data itself. But this idea didn't make a consensus at the time

[deanm0000]

I was originally going to make partition_by act like the existing group_by but from the linked issue it seems that the consensus is to treat ["a"] as "a" so I made the change to the GroupBy class to check for group count directly instead of based on input types.

I'll close this if/when the PR is accepted

[mcrumiller]

I agree with the OP here: treat "a" as a scalar and ["a"] as a tuple. Otherwise this requires people to add extra conditional layers to their code, with something like:

cols = get_partition_columns()  # returns list, we don't know the length
out = df.partition_by(columns, as_dict=True)

# now we have to deal with the fact that the list might have been length 1
if len(columns) == 1:
    out = {(k,): v for (k,v) in out.items()

[CaselIT]

@mcrumiller that's exactly my use case :)

[CaselIT]

This behaviour seems maintained also when using selectors, at least according to the docs: https://docs.pola.rs/py-polars/html/reference/dataframe/api/polars.DataFrame.partition_by.html last example df.partition_by(cs.string(), as_dict=True)

This means that depending on the number of columns matched by the selector a tuple or a scalar is returned. I think this is a very poor api behaviour, but this may be a different issue

[cmdlineluser]

I was going to ask if .partition_by could just be implemented in terms of .group_by so they behave the same?

But experimenting with your latest selectors example, there seems to be an issue with .group_by also

>>> [k for k, _ in df.group_by("a", "b", maintain_order=True)]
[(1, 0), (1, 1), (2, 2), (2, 3)]

>>> [k for k, _ in df.group_by(cs.all(), maintain_order=True)]
[1, 1, 2, 2] # <- Hmm?

[deanm0000]

@cmdlineluser It's because of

https://github.com/pola-rs/polars/blob/7ff3bf7b65396a20d958192bd2a36edc139e64ff/py-polars/polars/dataframe/group_by.py#L107-L111

My PR fixes that by counting the actual columns rather than inferring that when by is a str or Expr that it means there's only 1 result. The good news is that the issue seems to only exists in __iter__.

[stinodego]

If by is a single string ~or expression~, the tuples should have a non-tuple key. In all other instances, it should be a tuple.

Same for the groupby __iter__.

[CaselIT]

If by is a single string or expression, the tuples should have a non-tuple key. In all other instances, it should be a tuple.

do you think adding a flag to force the use of tuples would make sense? from an user prospective it's a lot nicer an api that has a stable return type

[stinodego]

I don't think we should add a parameter for this. If the current behavior is considered problematic, we should always use tuples to avoid the potential ambiguity.

[mcrumiller]

I don't think we should add a parameter for this. If the current behavior is considered problematic, we should always use tuples to avoid the potential ambiguity.

Agree. I don't think it's a lot nicer from a user perspective: the user can simply supply a tuple in if they want a tuple out. I don't necessarily consider the packaging of a dtype to necessary be a change in the dtype. In numpy, adding 1 to a scalar returns a scalar, adding 1 to a 2D array returns a 2D array, etc. In this example, we are partitioning by a scalar versus a list, and in the former case we return scalar keys, and in the latter list keys.

[CaselIT]

I may have misread @stinodego initial comment. Is the proposal that single string returns a scalar while a list/tuple of a single element returns a tuple of one element? if so that works for me (it was the original issue), sorry if I misunderstood.

(this leaves out the use of a single selector like cs.strings() that may match one or more columns depending on the dataframe, but I would selfishly be fine with ignoring this case it since it's not my use case now :))

[stinodego]

(this leaves out the use of a single selector like cs.strings() that may match one or more columns depending on the dataframe, but I would selfishly be fine with ignoring this case it since it's not my use case now :))

A single selector may match multiple columns so it has to be a tuple key. In that sense it's the same as a sequence input.

[mcrumiller]

Is the proposal that single string returns a scalar while a list/tuple of a single element returns a tuple of one element?

yes.

this leaves out the use of a single selector like cs.strings() that may match one or more columns

This is a good point. The cs module is entirely on the python side and so I think that we could detect when selectors are used and always return a tuple for those cases, although it is possible that some selectors are guaranteed to return single columns. I am not sure whether we should try to "push down" the scalar-vs-column logic into selectors, as that seems non-obvious.

[CaselIT]

I am not sure whether we should try to "push down" the scalar-vs-column logic into selectors, as that seems non-obvious.

agreed the fact that a scalar selector (meaning a single one, not a list of them) may match more than a column complicate things, but this is likely another issue

[myuanz]

I recently received a warning:

`group_by` iteration will change to always return group identifiers as tuples. Pass `by` as a list to silence this warning.

However, I believe that maintaining the current state, where different types of 'by' return different keys, is quite beneficial. As mentioned:

a single string returns a scalar while a list/tuple of a single element returns a tuple of one element.

This approach offers clear and intuitive semantics.

Following the update in #13646 and modifying the code as per the warning, besides the previously discussed semantic concern, the new syntax seems somewhat verbose, requiring an additional pair of parentheses and a comma:

- for k, g in df.group_by(....)
+ for (k, ), g in df.group_by(....)

Often, I use for in df.group_by instead of df.group_by().agg() merely for debugging purposes, and revert to the agg syntax once debugging is complete.

Therefore, it would be greatly appreciated if this warning could be removed and the deprecation halted as is.

Thank you for this wonderful tool, which I use daily.

[stinodego]

This approach offers clear and intuitive semantics.

I don't think it does. It is clear if you read the docs, but you might try group_by(pl.col("a")) and find that you get a tuple instead of a non-tuple key. Or you might try pl.first() and again find that while grouping by a single column, you get a tuple result.

The fact that group_by("a") behaves differently than group_by(col("a")) is the opposite of clear and intuitive. We've had multiple issues advocating for various behaviors. If we always return tuples, there is no more confusion.

Regarding your use case: it might indeed be a bit more verbose to write that way. But depending on what you do with the group key, maybe just using the one-tuple is also fine and you may not need to extract it? Or use k[0] whenever you do use it? The benefit is that if you decide to change the group by call, the type of k remains the same.

[myuanz]

So, when the deprecation of the single key feature in group_by eventually takes place, what will the parameter of group_by be? The options are:

1. *by: IntoExpr
2. by: Iterable[IntoExpr]
3. by: IntoExpr | Iterable[IntoExpr], *more_by: IntoExpr

Unless it's the second scenario, users might be compelled to change a single group_by parameter to (key, ) to eliminate warnings. However, if we look at the parameters for functions like with_columns, select, agg, and others, they all follow the first scenario. Therefore, no matter which signature is adopted after the deprecation, some users will be injured:

1. If options 1 or 3 are chosen, some users have to be warned until the deprecation is complete.
2. If option 2 is chosen, it will lead to inconsistency with the parameter conventions of existing systems, causing potential confusion.

How will things end up?

[stinodego]

The group_by function signature has been updated with 0.20.7. You can check out the docs.

[myuanz]

The signature for option 3 is copied from the source code. If you mean to say that the signature will remain the same even after the deprecation is complete, the following would be equivalent in the future:

df.group_by(["a"])
df.group_by("a")

However, before the deprecation is fully implemented, users who use df.group_by("a") and prefer to receive tuple returns would continually face warnings.

Regardless, I respect the decision of the development team. After all, having reached this point, any resolution might seem a bit peculiar.

[stinodego]

the following would be equivalent in the future

That's correct.

Unfortunately, single string input will warn for now. That's the deal with deprecation warnings.

[char101]

I just got bitten by this warning

DeprecationWarning: `partition_by(..., as_dict=True)` will change to always return tuples as dictionary keys. Pass `by` as a list to silence this warnin g, e.g. `partition_by(['key'], as_dict=True).

I actually didn't read the warning too much. What I thought was that the function signature has changed to only accept list, so I did that, and the code broke, because now the partition has become a tuple. I think that pass by notice should be removed or added and you should change your code to process the key as tuple.

Also, I don't think making the partition key tuple only make sense.

Group by in SQL or doing group by manually using dict in python always use the type of the key of the grouping type. It does not make sense that I have to type SELECT (age,), COUNT(1) FROM table GROUP BY (age, ) when grouping by a single key.

People that want uniform key type should also use uniform grouping key type, i.e. when grouping by a single column use a tuple with one element, when grouping by two columns use a tuple with two elements. So the responsibility is put on the user itself.

The library itself should not force the result type of the key.

[deanm0000]

@stinodego What do you think about making a different warning class so that it is easier to selectively ignore?

I'm thinking about putting (something like this) in polars exceptions

class PartitionByDeprecationWarning(DeprecationWarning):
    """Warning issued when partition_by used with as_dict and single key"""

and then changing the partition_by warning to that class.

and then people can do

warnings.filterwarnings("ignore", category=PartitionByDeprecationWarning)

I'm struggling to get the regex right for the message parameter of filterwarnings.

[stinodego]

I don't think that's a good idea. People running into this warning should ideally address it rather than ignore it. If they really want to ignore it, they can just ignore all deprecation warnings.

[s-banach]

Is the 1.0 behavior really going to be singleton tuples as keys for df.partition_by("a", as_dict=True)? It seems like the only reasoning is that some users may write pl.col("a") instead of "a".

If the underlying issue is that pl.col("a") is expected to behave identically to "a", you should address the underlying issue by making polars translate pl.col("a") to "a" under the hood. Changing the behavior of group_by and partition_by is just treating a symptom.

[char101]

The argument to partition_by is ambiguous between column names, expr and selector.

Maybe we could have a partition_by_col function which only accepts column names as string. If two or more column names are given then the dictionary is nested.