[Python] Joining two polars dataframes via a matching key and range overlaps on specified...

Suppose I have two dataframes generated with the code below:

df = pl.DataFrame({
'ID':[1, 2, 1, 2, 1, 2, 1, 1, 2, 2],
'Reference Value':[1.2, 3.5, 2.2, 1.4, 1.5, 3.7, 1.9, 1.1, 2.7, 3.8]
})

df2 = pl.DataFrame({
'ID':[1, 1, 1, 2, 2, 2, 3, 3, 3],
'Lower Reference Range':[1, 2, 3, 1, 2, 3, 1, 2, 3],
'Upper Reference Range':[2, 3, 4, 2, 3, 4, 2, 3, 4],
'Map Value':['a1','b1','c1','a2','b2', 'c2', 'a3', 'b3', 'c3']
})


Which generates a dataframe as such:

# df
shape: (10, 2)
┌─────┬─────────────────┐
│ ID ┆ Reference Value │
│ --- ┆ --- │
│ i64 ┆ f64 │
╞═════╪═════════════════╡
│ 1 ┆ 1.2 │
│ 2 ┆ 3.5 │
│ 1 ┆ 2.2 │
│ 2 ┆ 1.4 │
│ 1 ┆ 1.5 │
│ 2 ┆ 3.7 │
│ 1 ┆ 1.9 │
│ 1 ┆ 1.1 │
│ 2 ┆ 2.7 │
│ 2 ┆ 3.8 │
└─────┴─────────────────┘
# df2
shape: (9, 4)
┌─────┬───────────────────────┬───────────────────────┬───────────┐
│ ID ┆ Lower Reference Range ┆ Upper Reference Range ┆ Map Value │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ str │
╞═════╪═══════════════════════╪═══════════════════════╪═══════════╡
│ 1 ┆ 1 ┆ 2 ┆ a1 │
│ 1 ┆ 2 ┆ 3 ┆ b1 │
│ 1 ┆ 3 ┆ 4 ┆ c1 │
│ 2 ┆ 1 ┆ 2 ┆ a2 │
│ 2 ┆ 2 ┆ 3 ┆ b2 │
│ 2 ┆ 3 ┆ 4 ┆ c2 │
│ 3 ┆ 1 ┆ 2 ┆ a3 │
│ 3 ┆ 2 ┆ 3 ┆ b3 │
│ 3 ┆ 3 ┆ 4 ┆ c3 │
└─────┴───────────────────────┴───────────────────────┴───────────┘



I would like to map df1 to df2 such that where the IDs match up and the value of Reference Value is between the lower and upper range in the reference table, the map value is assigned. As below:

┌─────┬─────────────────┬───────────┐
│ ID ┆ Reference Value ┆ Map Value │
│ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ str │
╞═════╪═════════════════╪═══════════╡
│ 1 ┆ 1.2 ┆ a1 │
│ 2 ┆ 3.5 ┆ c2 │
│ 1 ┆ 2.2 ┆ b1 │
│ 2 ┆ 1.4 ┆ a2 │
│ ... ┆ ... ┆ ... │
│ 1 ┆ 1.9 ┆ a1 │
│ 1 ┆ 1.1 ┆ a1 │
│ 2 ┆ 2.7 ┆ b2 │
│ 2 ┆ 3.8 ┆ c2 │
└─────┴─────────────────┴───────────┘


I've implemented a simple naïve approach via the when and then method as below:

def simple_when_then_mapper(df,df2,ref_id,ref_value,lower_range_col,upper_range_col,mapping_col):
'''
df = dataframe to map values unto
df2: reference dataframe
ref_id, ref_value,lower_range_col,upper_range_col,mapping_col: str referencing the column containing the Reference ID, Reference Value, etc
'''
for i in range(df2.shape[0]):
if 'query' not in locals():
query=(pl.when((pl.col(ref_id)==df2[ref_id]) &
(pl.col(ref_value)>df2[lower_range_col])&
(pl.col(ref_value)<=df2[upper_range_col])
)
.then(df2[mapping_col]))
else:
query=(query.when((pl.col(ref_id)==df2[ref_id]) &
(pl.col(ref_value)>df2[lower_range_col])&
(pl.col(ref_value)<=df2[upper_range_col])
)
.then(df2[mapping_col]))
query=query.otherwise(None).alias(mapping_col)
return df.with_columns(query)


This works but my concern is when the tables and cardinality are sufficiently high, I'll run into issues. From my testing, when-then nesting of ~400 starts to trigger stack overflows on my machine. Is there a more proper and scalable way to do the above?

Continue reading...