Support mask= argument in LocBody

[jrycw]

Related Issue: #389.

Hello team,

This is a rough idea for supporting the passing of Polars expressions to the mask= argument of LobBody.

The procedures are outlined as follows:

1. Exclude the stub and group columns if necessary.
2. Use df.select() to process the Polars expression and obtain a boolean-like matrix.
3. Use two for-loops to iterate through the cells and select the desired cells based on whether their values are True.

This idea is preliminary but should provide a foundation for further discussions.

import polars as pl
from great_tables import GT, loc, style
from great_tables.data import gtcars
from polars import selectors as cs

df_mini = pl.from_pandas(gtcars).head().pivot(on="year", index=["trim", "mfr"], values="hp")
(
    GT(df_mini,rowname_col="trim", groupname_col="mfr")
    .tab_style(
        style=style.text(color="red"),
        locations=loc.body(mask=cs.numeric().gt(600))
    )
)

print(df_mini.select(cs.numeric().gt(600)))

shape: (3, 4)
┌────────┬────────┬────────┬────────┐
│ 2017.0 ┆ 2015.0 ┆ 2014.0 ┆ 2016.0 │
│ ---    ┆ ---    ┆ ---    ┆ ---    │
│ bool   ┆ bool   ┆ bool   ┆ bool   │
╞════════╪════════╪════════╪════════╡
│ true   ┆ null   ┆ null   ┆ null   │
│ null   ┆ false  ┆ false  ┆ true   │
│ null   ┆ false  ┆ null   ┆ null   │
└────────┴────────┴────────┴────────┘

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.73%. Comparing base (4043f14) to head (be7c141).
Report is 15 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #566      +/-   ##
==========================================
+ Coverage   90.67%   90.73%   +0.06%     
==========================================
  Files          46       46              
  Lines        5381     5416      +35     
==========================================
+ Hits         4879     4914      +35     
  Misses        502      502              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[machow]

What if the rules for the mask were as follows?:

• the mask expressions should--when run by frame.select(mask_expr)--return a DataFrame that:
  • columns: has the same or fewer columns (names must match original frame, but could be in a different order).
  • rows: has the same number of rows
• after running the expression, we validate this right away
• the mask is assumed to be in the original row order
• in this case we can capture the column name and row number of each cell (e.g. with enumerate(mask_result.iter_rows())).

[jrycw]

Thanks for summarizing the rules.

It seems that for the last part, we only need a single loop to gather all the information required by CellPos.

[jrycw]

How about this version?

For prototyping, use assert for validation, which can later be replaced with raise ValueError().

[jrycw]

I've replaced assert with raise ValueError() and added more test cases.

[machow]

@jrycw this is looking really good! Anything else needed for this PR? I'm happy to help with any pieces!

[jrycw]

@machow One thing I’m unsure about is the naming of the function. Specifically, I’m not clear on what the i or c represents in resolve_rows_i and resolve_cols_c. Does resolve_mask_i make sense? Other than that, I think the feature is ready for users to experiment with. The idea of mask= is great, and I believe it will offer a lot of flexibility to users.

[rich-iannone]

@machow One thing I’m unsure about is the naming of the function. Specifically, I’m not clear on what the i or c represents in resolve_rows_i and resolve_cols_c. Does resolve_mask_i make sense? Other than that, I think the feature is ready for users to experiment with. The idea of mask= is great, and I believe it will offer a lot of flexibility to users.

I can help!

i == index (int)
c == column (str)

this is admittedly a bit of a holdover from the R package (which also has the resolve_cols_i() and resolve_rows_l() functions). But based on this, maybe just resolve_mask is fine?

[jrycw]

@rich-iannone Thanks for the explanation! I agree that resolve_mask seems like a better fit here.

[machow]

From a quick glance it's looking really nice -- if it's ready for review I'm happy to take a closer look!

[jrycw]

The PR should be ready for review; please proceed.

[machow]

This is looking really great--thanks for taking the time to implement this! I left several small suggestions for some error wording, and tweaks. Always happy to pick things up from here if useful. Thanks again for all the work you put into this!

[jrycw]

@machow, I truly appreciate your review—all your suggestions are spot on! Also, it’s great to see all green checks again.

[machow]

This looks really great, thanks! @rich-iannone do you want to merge?