How to replace a conditional for loop with a more efficient approach using Polars?

SB_dates = (
   SB_dates
    .with_row_index("row")
    .unpivot(index="row", variable_name="col", value_name="date")
    .with_columns(pl.col("col").rle_id())
)

df_mask = pl.concat(
   [SB_dates, newBL.unpivot().select(is_block = "value")], 
   how = "horizontal"
)

df_mask = df_mask.with_columns(
   pl.when(pl.col.is_block)
     .then(pl.col.date)
     .forward_fill()
     .alias("date_block")
)

all_ids = data.group_by(pl.col.date.dt.date()).agg("ids")

random_ids = (
   data
    .group_by(pl.col.date.dt.date())
    .agg(random_ids = pl.col.ids.shuffle().head(5))
)

df_selections = (
   df_mask
    .join(all_ids, on="date")
    .join(random_ids, left_on="date_block", right_on="date")
    .with_columns(pl.col.ids.list.set_intersection(pl.col.random_ids))
    .select("row", "col", "date", "ids")
    .explode("ids")
)

result = (
   df_selections
    .join(data, on=[pl.col.date.dt.date(), "ids"])
    .sort("row", "col")
    .pivot(
       on = "col", 
       index = "row", 
       values = "values", 
       aggregate_function = pl.element().mean()
    )
)

Reshaping

Generally, if you can reshape your data from "wide" to "long", or turn "columns" into "rows" - this is where Polars will be most efficient.

.unpivot()

SB_dates = (
   SB_dates
    .with_row_index("row")
    .unpivot(index="row", variable_name="col", value_name="date")
    .with_columns(pl.col("col").rle_id())
)

shape: (120, 3)
┌─────┬─────┬────────────┐
│ row ┆ col ┆ date       │
│ --- ┆ --- ┆ ---        │
│ u32 ┆ u32 ┆ date       │
╞═════╪═════╪════════════╡
│ 0   ┆ 0   ┆ 2009-03-31 │
│ 1   ┆ 0   ┆ 2010-05-31 │
│ 2   ┆ 0   ┆ 2011-03-31 │
│ 3   ┆ 0   ┆ 2001-12-31 │
│ 4   ┆ 0   ┆ 2000-06-30 │
│ …   ┆ …   ┆ …          │
│ 15  ┆ 5   ┆ 2020-11-30 │
│ 16  ┆ 5   ┆ 2012-05-31 │
│ 17  ┆ 5   ┆ 2011-02-28 │
│ 18  ┆ 5   ┆ 2008-03-31 │
│ 19  ┆ 5   ┆ 2023-07-31 │
└─────┴─────┴────────────┘

• .with_row_index() and .rle_id() are used to to assign row, col numbers which will be used to recreate the original shape later.

We also reshape newBL and horizontally concat to pair up the corresponding boolean values.

This is used to then mask out and .forward_fill() the date column.

df_mask = pl.concat([SB_dates, newBL.unpivot().select(is_block = "value")], how="horizontal")

df_mask = df_mask.with_columns(
   pl.when(pl.col.is_block)
     .then(pl.col.date)
     .forward_fill()
     .alias("date_block")
)

shape: (120, 5)
┌─────┬─────┬────────────┬──────────┬────────────┐
│ row ┆ col ┆ date       ┆ is_block ┆ date_block │
│ --- ┆ --- ┆ ---        ┆ ---      ┆ ---        │
│ u32 ┆ u32 ┆ date       ┆ bool     ┆ date       │
╞═════╪═════╪════════════╪══════════╪════════════╡
│ 0   ┆ 0   ┆ 2009-03-31 ┆ true     ┆ 2009-03-31 │ # <-
│ 1   ┆ 0   ┆ 2010-05-31 ┆ false    ┆ 2009-03-31 │ # <-
│ 2   ┆ 0   ┆ 2011-03-31 ┆ false    ┆ 2009-03-31 │ # <-
│ 3   ┆ 0   ┆ 2001-12-31 ┆ false    ┆ 2009-03-31 │ # <-
│ 4   ┆ 0   ┆ 2000-06-30 ┆ true     ┆ 2000-06-30 │
│ …   ┆ …   ┆ …          ┆ …        ┆ …          │
│ 15  ┆ 5   ┆ 2020-11-30 ┆ false    ┆ 2010-05-31 │
│ 16  ┆ 5   ┆ 2012-05-31 ┆ true     ┆ 2012-05-31 │
│ 17  ┆ 5   ┆ 2011-02-28 ┆ true     ┆ 2011-02-28 │
│ 18  ┆ 5   ┆ 2008-03-31 ┆ true     ┆ 2008-03-31 │
│ 19  ┆ 5   ┆ 2023-07-31 ┆ true     ┆ 2023-07-31 │
└─────┴─────┴────────────┴──────────┴────────────┘

By using date_block to search for ids, we will emulate the desired behaviour:

If it's not a new block, I should retain the same random ids that were previously selected.

List is_in() other_list

If we reshape data so that we have lists of ids and random_ids

all_ids = data.group_by(pl.col.date.dt.date()).agg("ids")

random_ids = (
   data
    .group_by(pl.col.date.dt.date())
    .agg(random_ids = pl.col.ids.shuffle().head(5))
)

.list.set_intersection() can be used to do the same as the mulitple .filter + .is_in checks you're using.

(
   df_mask
    .join(all_ids, on="date")
    .join(random_ids, left_on="date_block", right_on="date")
    .with_columns(selection = pl.col.ids.list.set_intersection(pl.col.random_ids))
)

shape: (114, 8)
┌─────┬─────┬────────────┬──────────┬────────────┬────────────────┬────────────────┬────────────────┐
│ row ┆ col ┆ date       ┆ is_block ┆ date_block ┆ ids            ┆ random_ids     ┆ selection      │
│ --- ┆ --- ┆ ---        ┆ ---      ┆ ---        ┆ ---            ┆ ---            ┆ ---            │
│ u32 ┆ u32 ┆ date       ┆ bool     ┆ date       ┆ list[i64]      ┆ list[i64]      ┆ list[i64]      │
╞═════╪═════╪════════════╪══════════╪════════════╪════════════════╪════════════════╪════════════════╡
│ 19  ┆ 3   ┆ 2005-10-31 ┆ false    ┆ 2019-01-31 ┆ [73, 38]       ┆ [80]           ┆ []             │
│ 10  ┆ 0   ┆ 2015-12-31 ┆ true     ┆ 2015-12-31 ┆ [97, 57, … 12] ┆ [57, 40, … 81] ┆ [97, 57, … 12] │
│ 4   ┆ 5   ┆ 2011-07-31 ┆ false    ┆ 2012-01-31 ┆ [98, 40, 31]   ┆ [48, 6]        ┆ []             │
│ 2   ┆ 5   ┆ 2012-01-31 ┆ true     ┆ 2012-01-31 ┆ [48, 6]        ┆ [48, 6]        ┆ [48, 6]        │
│ 16  ┆ 3   ┆ 2010-08-31 ┆ true     ┆ 2010-08-31 ┆ [35, 15, … 54] ┆ [15, 19, … 35] ┆ [35, 15, … 54] │
│ …   ┆ …   ┆ …          ┆ …        ┆ …          ┆ …              ┆ …              ┆ …              │
│ 5   ┆ 2   ┆ 2009-11-30 ┆ false    ┆ 2018-12-31 ┆ [92, 52, … 58] ┆ [50, 33, … 97] ┆ []             │
│ 6   ┆ 5   ┆ 2019-03-31 ┆ true     ┆ 2019-03-31 ┆ [88, 12, … 2]  ┆ [35, 2, … 88]  ┆ [88, 12, … 2]  │
│ 9   ┆ 0   ┆ 2021-04-30 ┆ false    ┆ 2020-08-31 ┆ [19]           ┆ [74, 34, … 87] ┆ []             │
│ 16  ┆ 1   ┆ 2017-09-30 ┆ true     ┆ 2017-09-30 ┆ [83, 93, … 13] ┆ [93, 66, … 59] ┆ [83, 93, … 13] │
│ 1   ┆ 5   ┆ 2010-10-31 ┆ false    ┆ 2010-09-30 ┆ [90, 13]       ┆ [75, 45]       ┆ []             │
└─────┴─────┴────────────┴──────────┴────────────┴────────────────┴────────────────┴────────────────┘

From here, we keep only the needed columns and .explode() the list back into rows.

df_selections = (
   df_mask
    .join(all_ids, on="date")
    .join(random_ids, left_on="date_block", right_on="date")
    .with_columns(pl.col.ids.list.set_intersection(pl.col.random_ids))
    .select("row", "col", "date", "ids")
    .explode("ids")
)

shape: (258, 4)
┌─────┬─────┬────────────┬──────┐
│ row ┆ col ┆ date       ┆ ids  │
│ --- ┆ --- ┆ ---        ┆ ---  │
│ u32 ┆ u32 ┆ date       ┆ i64  │
╞═════╪═════╪════════════╪══════╡
│ 0   ┆ 4   ┆ 2006-01-31 ┆ 60   │
│ 0   ┆ 4   ┆ 2006-01-31 ┆ 7    │
│ 0   ┆ 4   ┆ 2006-01-31 ┆ 70   │
│ 0   ┆ 4   ┆ 2006-01-31 ┆ 62   │
│ 0   ┆ 4   ┆ 2006-01-31 ┆ 82   │
│ …   ┆ …   ┆ …          ┆ …    │
│ 9   ┆ 4   ┆ 2015-03-31 ┆ null │
│ 8   ┆ 2   ┆ 2008-06-30 ┆ 70   │
│ 8   ┆ 2   ┆ 2008-06-30 ┆ 12   │
│ 19  ┆ 5   ┆ 2023-07-31 ┆ 12   │
│ 19  ┆ 5   ┆ 2023-07-31 ┆ 67   │
└─────┴─────┴────────────┴──────┘

Pivot

With only the wanted ids remaining, a .join() can be used to get the corresponding values.

(df_selections
  .join(data, on=[pl.col.date.dt.date(), "ids"])
  .sort("row", "col")
)

shape: (261, 9)
┌─────┬─────┬────────────┬──────────┬───┬──────┬─────────────────────┬───────────┬───────────┐
│ row ┆ col ┆ date       ┆ is_block ┆ … ┆ ids  ┆ date_right          ┆ values    ┆ ids_right │
│ --- ┆ --- ┆ ---        ┆ ---      ┆   ┆ ---  ┆ ---                 ┆ ---       ┆ ---       │
│ u32 ┆ u32 ┆ date       ┆ bool     ┆   ┆ i64  ┆ datetime[μs]        ┆ f64       ┆ i64       │
╞═════╪═════╪════════════╪══════════╪═══╪══════╪═════════════════════╪═══════════╪═══════════╡
│ 0   ┆ 0   ┆ 2009-03-31 ┆ true     ┆ … ┆ 24   ┆ 2009-03-31 00:00:00 ┆ -0.006142 ┆ 24        │
│ 0   ┆ 1   ┆ 2014-06-30 ┆ true     ┆ … ┆ 27   ┆ 2014-06-30 00:00:00 ┆ -0.001445 ┆ 27        │
│ 0   ┆ 1   ┆ 2014-06-30 ┆ true     ┆ … ┆ 57   ┆ 2014-06-30 00:00:00 ┆ 0.005532  ┆ 57        │
│ 0   ┆ 1   ┆ 2014-06-30 ┆ true     ┆ … ┆ 96   ┆ 2014-06-30 00:00:00 ┆ -0.00793  ┆ 96        │
│ 0   ┆ 1   ┆ 2014-06-30 ┆ true     ┆ … ┆ 14   ┆ 2014-06-30 00:00:00 ┆ 0.003196  ┆ 14        │
│ …   ┆ …   ┆ …          ┆ …        ┆ … ┆ …    ┆ …                   ┆ …         ┆ …         │
│ 19  ┆ 2   ┆ 2021-10-31 ┆ true     ┆ … ┆ 15   ┆ 2021-10-31 00:00:00 ┆ 0.003261  ┆ 15        │
│ 19  ┆ 3   ┆ 2005-10-31 ┆ false    ┆ … ┆ null ┆ null                ┆ null      ┆ null      │
│ 19  ┆ 4   ┆ 2000-03-31 ┆ false    ┆ … ┆ null ┆ null                ┆ null      ┆ null      │
│ 19  ┆ 5   ┆ 2023-07-31 ┆ true     ┆ … ┆ 12   ┆ 2023-07-31 00:00:00 ┆ -0.000579 ┆ 12        │
│ 19  ┆ 5   ┆ 2023-07-31 ┆ true     ┆ … ┆ 67   ┆ 2023-07-31 00:00:00 ┆ -0.009836 ┆ 67        │
└─────┴─────┴────────────┴──────────┴───┴──────┴─────────────────────┴───────────┴───────────┘

.pivot() can then be used to return back to the original "wide" shape and aggregate the mean.

(df_selections
  .join(data, on=[pl.col.date.dt.date(), "ids"])
  .sort("row", "col")
  .pivot(on="col", index="row", values="values", aggregate_function=pl.element().mean())
)

shape: (20, 7)
┌─────┬───────────┬───────────┬───────────┬───────────┬───────────┬───────────┐
│ row ┆ 0         ┆ 1         ┆ 2         ┆ 3         ┆ 4         ┆ 5         │
│ --- ┆ ---       ┆ ---       ┆ ---       ┆ ---       ┆ ---       ┆ ---       │
│ u32 ┆ f64       ┆ f64       ┆ f64       ┆ f64       ┆ f64       ┆ f64       │
╞═════╪═══════════╪═══════════╪═══════════╪═══════════╪═══════════╪═══════════╡
│ 0   ┆ -0.006142 ┆ -0.000162 ┆ 0.001333  ┆ 0.000646  ┆ 0.004068  ┆ -0.007742 │
│ 1   ┆ null      ┆ -0.000778 ┆ null      ┆ null      ┆ null      ┆ null      │
│ 2   ┆ null      ┆ null      ┆ null      ┆ null      ┆ -0.001934 ┆ -0.004242 │
│ 3   ┆ null      ┆ null      ┆ -0.002209 ┆ -0.005019 ┆ null      ┆ 0.004758  │
│ 4   ┆ null      ┆ null      ┆ -0.000181 ┆ 0.001981  ┆ null      ┆ null      │
│ …   ┆ …         ┆ …         ┆ …         ┆ …         ┆ …         ┆ …         │
│ 15  ┆ null      ┆ -0.00059  ┆ 0.003549  ┆ null      ┆ null      ┆ null      │
│ 16  ┆ null      ┆ -0.003033 ┆ -0.003727 ┆ -0.002882 ┆ null      ┆ 0.009226  │
│ 17  ┆ null      ┆ null      ┆ null      ┆ -0.003837 ┆ null      ┆ 0.009743  │
│ 18  ┆ null      ┆ -0.001557 ┆ 0.000865  ┆ -0.008088 ┆ null      ┆ 0.001735  │
│ 19  ┆ 0.008602  ┆ 0.000155  ┆ -0.001692 ┆ null      ┆ null      ┆ -0.005207 │
└─────┴───────────┴───────────┴───────────┴───────────┴───────────┴───────────┘