sort_by gives incorrect results when preceded by filter in an aggregation context

[bdholder]

Checks

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

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

Reproducible example

Code

import polars as pl

df = pl.DataFrame({
    "id": [1, 1, 1, 2, 2, 3, 3, 3],
    "number": [1, 3, 2, 1, 2, 2, 1, 3],
    "type": ["A", "B", "A", "B", "B", "A", "B", "C"],
    "cost": [10, 25, 20, 25, 30, 30, 50, 100],
})
print(df)

cost = pl.col("cost")
number = pl.col("number")
type = pl.col("type")

sort_by_filter = cost.sort_by(number).filter(type == "A")
filter_sort_by = cost.filter(type == "A").sort_by(number)

agg_df = df.group_by("id").agg(
    sort_by_first=sort_by_filter.first(),
    raw_sort_by_first=sort_by_filter,
    filter_first=filter_sort_by.first(),
    raw_filter_first=filter_sort_by,
)
print(agg_df)

select_df = df.select(
    sort_by_first=cost.sort_by(number).filter(type == "A").first(),
    raw_sort_by_first=cost.sort_by(number).filter(type == "A").implode(),
)
print(select_df)

bad_df = df.select(
    filter_first=filter_sort_by.first(),
)

Output

shape: (8, 4)
┌─────┬────────┬──────┬──────┐
│ id  ┆ number ┆ type ┆ cost │
│ --- ┆ ---    ┆ ---  ┆ ---  │
│ i64 ┆ i64    ┆ str  ┆ i64  │
╞═════╪════════╪══════╪══════╡
│ 1   ┆ 1      ┆ A    ┆ 10   │
│ 1   ┆ 3      ┆ B    ┆ 25   │
│ 1   ┆ 2      ┆ A    ┆ 20   │
│ 2   ┆ 1      ┆ B    ┆ 25   │
│ 2   ┆ 2      ┆ B    ┆ 30   │
│ 3   ┆ 2      ┆ A    ┆ 30   │
│ 3   ┆ 1      ┆ B    ┆ 50   │
│ 3   ┆ 3      ┆ C    ┆ 100  │
└─────┴────────┴──────┴──────┘
shape: (3, 5)
┌─────┬───────────────┬───────────────────┬──────────────┬──────────────────┐
│ id  ┆ sort_by_first ┆ raw_sort_by_first ┆ filter_first ┆ raw_filter_first │
│ --- ┆ ---           ┆ ---               ┆ ---          ┆ ---              │
│ i64 ┆ i64           ┆ list[i64]         ┆ i64          ┆ list[i64]        │
╞═════╪═══════════════╪═══════════════════╪══════════════╪══════════════════╡
│ 2   ┆ null          ┆ []                ┆ 25           ┆ [25, 30]         │
│ 1   ┆ 10            ┆ [10, 20]          ┆ 10           ┆ [10, 20, 25]     │
│ 3   ┆ 30            ┆ [30]              ┆ 50           ┆ [50, 30, 100]    │
└─────┴───────────────┴───────────────────┴──────────────┴──────────────────┘
shape: (1, 2)
┌───────────────┬───────────────────┐
│ sort_by_first ┆ raw_sort_by_first │
│ ---           ┆ ---               │
│ i64           ┆ list[i64]         │
╞═══════════════╪═══════════════════╡
│ 10            ┆ [10, 50, 30]      │
└───────────────┴───────────────────┘

Log output

keys/aggregates are not partitionable: running default HASH AGGREGATION
Traceback (most recent call last):
  File "C:\Users\bholder\code\notebooks\bug_test.py", line 32, in <module>
    bad_df = df.select(
  File "C:\Users\bholder\code\notebooks\.venv\lib\site-packages\polars\dataframe\frame.py", line 7727, in select
    return self.lazy().select(*exprs, **named_exprs).collect(eager=True)
  File "C:\Users\bholder\code\notebooks\.venv\lib\site-packages\polars\utils\deprecation.py", line 95, in wrapper
    return function(*args, **kwargs)
  File "C:\Users\bholder\code\notebooks\.venv\lib\site-packages\polars\lazyframe\frame.py", line 1711, in collect
    return wrap_df(ldf.collect())
exceptions.ComputeError: `sort_by` produced different length: 8 than the series that has to be sorted: 3

Error originated in expression: 'col("cost").filter([(col("type")) == (Utf8(A))]).sort_by(by=[col("number")], descending=[false])'

Issue description

The behavior of sort_by is inconsistent. I'm not entirely certain about the correct mental model for understanding "multi-column" operations. In particular, I'm a bit surprised that the aggregate "sort_by_first" column gives the correct result.

Expected behavior

I expected sort_by to behave similarly in an aggregation context and in a general selection context, and especially for it to not "discard" a preceding filter operation.

Installed versions

--------Version info---------
Polars:              0.19.5
Index type:          UInt32
Platform:            Windows-10-10.0.19045-SP0
Python:              3.10.1 (tags/v3.10.1:2cd268a, Dec  6 2021, 19:10:37) [MSC v.1929 64 bit (AMD64)]

----Optional dependencies----
adbc_driver_sqlite:  <not installed>
cloudpickle:         <not installed>
connectorx:          0.3.2
deltalake:           <not installed>
fsspec:              <not installed>
gevent:              <not installed>
matplotlib:          3.8.0
numpy:               1.26.0
openpyxl:            <not installed>
pandas:              1.4.1
pyarrow:             13.0.0
pydantic:            <not installed>
pyiceberg:           <not installed>
pyxlsb:              <not installed>
sqlalchemy:          2.0.21
xlsx2csv:            <not installed>
xlsxwriter:          <not installed>

[cmdlineluser]

Just in trying to understand the issue better:

df = pl.DataFrame({
    "id": [1, 1, 1, 2, 2, 3, 3, 3],
    "number": [1, 3, 2, 1, 2, 2, 1, 3],
    "type": ["A", "B", "A", "B", "B", "A", "B", "C"],
    "cost": [10, 25, 20, 25, 30, 30, 50, 100],
})

df.group_by("id", maintain_order=True).agg(
   OK    = pl.col("cost").filter(pl.col("type") == "A"),
   ERROR = pl.col("cost").filter(pl.col("type") == "A").sort_by("number")
)

# shape: (3, 3)
# ┌─────┬───────────┬───────────────┐
# │ id  ┆ OK        ┆ ERROR         │
# │ --- ┆ ---       ┆ ---           │
# │ i64 ┆ list[i64] ┆ list[i64]     │
# ╞═════╪═══════════╪═══════════════╡
# │ 1   ┆ [10, 20]  ┆ [10, 20, 25]  │
# │ 2   ┆ []        ┆ [25, 30]      │
# │ 3   ┆ [30]      ┆ [50, 30, 100] │
# └─────┴───────────┴───────────────┘

Adding .sort_by seems to be adding extra values from elsewhere and this example should actually raise a ComputeError due to a length mismatch?

[mcrumiller]

It looks like the filter is being completely ignored:

df.group_by("id", maintain_order=True).agg(
    FILTERED = pl.col("cost").filter(pl.col("type") == "A").sort_by("number"),
    NOT_FILTERED = pl.col("cost").sort_by("number")
)

shape: (3, 3)
┌─────┬───────────────┬───────────────┐
│ id  ┆ FILTERED      ┆ NOT_FILTERED  │
│ --- ┆ ---           ┆ ---           │
│ i64 ┆ list[i64]     ┆ list[i64]     │
╞═════╪═══════════════╪═══════════════╡
│ 1   ┆ [10, 20, 25]  ┆ [10, 20, 25]  │
│ 2   ┆ [25, 30]      ┆ [25, 30]      │
│ 3   ┆ [50, 30, 100] ┆ [50, 30, 100] │
└─────┴───────────────┴───────────────┘

[ritchie46]

We should error. The two columns will have different number of groups.