Multi filter by 2 columns and display some best results with Polars
Question:
I have df for my work with 3 main columns: cid1, cid2, cid3, and more columns cid4, cid5, etc. cid1 and cid2 is int, another columns is float.
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
| 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
Each combitations of cid1 and cid2 is a workset for analysis and for each workset I have some values cid3.
I can take df with only maximal values of cid3:
df = df.filter(pl.col("cid3") == pl.col("cid3").max().over(["cid1", "cid2"]))
And this display:
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
But I can’t catch how I can take two maximal values of cid3 for each workset for this result:
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
| 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
Two maximal values of cid3 is for example, for my task I can take 10 maximal values, 5 minimal and so that. Help me please!
Answers:
You can use .top_k() to get the k largest (or smallest) values.
.unique().top_k() can be used if you need distinct values.
df.groupby("cid1", "cid2").agg(pl.col("cid3").top_k(2))
shape: (2, 3)
┌──────┬──────┬────────────┐
│ cid1 ┆ cid2 ┆ cid3 │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ list[f64] │
╞══════╪══════╪════════════╡
│ 1 ┆ 5 ┆ [9.0, 2.0] │
│ 3 ┆ 7 ┆ [8.0, 3.0] │
└──────┴──────┴────────────┘
This can be used inside .filter combined with .is_in
df.filter(
pl.col("cid3").is_in(pl.col("cid3").top_k(2))
.over("cid1", "cid2")
)
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
descending=True to find the minimal values (bottom_k)
Update: .bottom_k has been added and will be in the next release.
df.filter(
pl.col("cid3").is_in(pl.col("cid3").bottom_k(2)
.over("cid1", "cid2")
)
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
Dataframe used:
df = pl.read_csv(b"""
cid1,cid2,cid3,cid4,cid5,cid6
1,5,1.0,4.0,4.0,1.0
1,5,2.0,5.0,5.0,9.0
1,5,9.0,6.0,4.0,9.0
3,7,1.0,7.0,9.0,1.0
3,7,3.0,7.0,9.0,1.0
3,7,8.0,8.0,3.0,1.0
""")
Here is one more possibility in case you want to get maximum or minimum values
Getting 2 largest values
df.filter(
pl.col("cid3").is_in(pl.col("cid3").unique().sort(descending=True).head(2))
.over(["cid1", "cid2"])
)
# Result
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
Getting 2 smallest values
(df.filter(
pl.col("cid3").is_in(pl.col("cid3").unique().sort(descending=False).head(2))
.over(["cid1", "cid2"])
)
# Result
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
I have df for my work with 3 main columns: cid1, cid2, cid3, and more columns cid4, cid5, etc. cid1 and cid2 is int, another columns is float.
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
| 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
Each combitations of cid1 and cid2 is a workset for analysis and for each workset I have some values cid3.
I can take df with only maximal values of cid3:
df = df.filter(pl.col("cid3") == pl.col("cid3").max().over(["cid1", "cid2"]))
And this display:
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
But I can’t catch how I can take two maximal values of cid3 for each workset for this result:
┌──────┬──────┬──────┬──────┬──────┬──────────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
╞══════╪══════╪══════╪══════╪══════╪══════════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
| 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────────┘
Two maximal values of cid3 is for example, for my task I can take 10 maximal values, 5 minimal and so that. Help me please!
You can use .top_k() to get the k largest (or smallest) values.
.unique().top_k() can be used if you need distinct values.
df.groupby("cid1", "cid2").agg(pl.col("cid3").top_k(2))
shape: (2, 3)
┌──────┬──────┬────────────┐
│ cid1 ┆ cid2 ┆ cid3 │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ list[f64] │
╞══════╪══════╪════════════╡
│ 1 ┆ 5 ┆ [9.0, 2.0] │
│ 3 ┆ 7 ┆ [8.0, 3.0] │
└──────┴──────┴────────────┘
This can be used inside .filter combined with .is_in
df.filter(
pl.col("cid3").is_in(pl.col("cid3").top_k(2))
.over("cid1", "cid2")
)
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
descending=True to find the minimal values (bottom_k)
Update: .bottom_k has been added and will be in the next release.
df.filter(
pl.col("cid3").is_in(pl.col("cid3").bottom_k(2)
.over("cid1", "cid2")
)
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
Dataframe used:
df = pl.read_csv(b"""
cid1,cid2,cid3,cid4,cid5,cid6
1,5,1.0,4.0,4.0,1.0
1,5,2.0,5.0,5.0,9.0
1,5,9.0,6.0,4.0,9.0
3,7,1.0,7.0,9.0,1.0
3,7,3.0,7.0,9.0,1.0
3,7,8.0,8.0,3.0,1.0
""")
Here is one more possibility in case you want to get maximum or minimum values
Getting 2 largest values
df.filter(
pl.col("cid3").is_in(pl.col("cid3").unique().sort(descending=True).head(2))
.over(["cid1", "cid2"])
)
# Result
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 1 ┆ 5 ┆ 9.0 ┆ 6.0 ┆ 4.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 8.0 ┆ 8.0 ┆ 3.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘
Getting 2 smallest values
(df.filter(
pl.col("cid3").is_in(pl.col("cid3").unique().sort(descending=False).head(2))
.over(["cid1", "cid2"])
)
# Result
shape: (4, 6)
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1 ┆ 5 ┆ 1.0 ┆ 4.0 ┆ 4.0 ┆ 1.0 │
│ 1 ┆ 5 ┆ 2.0 ┆ 5.0 ┆ 5.0 ┆ 9.0 │
│ 3 ┆ 7 ┆ 1.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
│ 3 ┆ 7 ┆ 3.0 ┆ 7.0 ┆ 9.0 ┆ 1.0 │
└──────┴──────┴──────┴──────┴──────┴──────┘