Skip to content

Commit

Permalink
refactor: use replace over map_dict (#231)
Browse files Browse the repository at this point in the history
  • Loading branch information
baggiponte authored Jun 8, 2024
1 parent d50d4e0 commit 0d41c1f
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions functime/base/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ def _set_string_cache(df: pl.DataFrame):
# Reset categorical to string type
df = df.with_columns(pl.col(entity_col).cast(pl.Utf8))
df_new = df.with_columns(
pl.col(entity_col).map_dict(string_cache, return_dtype=pl.Int32)
pl.col(entity_col).replace(string_cache, return_dtype=pl.Int32, default=None)
)
inv_string_cache = {i: entity for entity, i in string_cache.items()}
return df_new, entity_col_dtype, string_cache, inv_string_cache
Expand All @@ -29,15 +29,17 @@ def _enforce_string_cache(
# Reset categorical to string type
df = df.with_columns(pl.col(entity_col).cast(pl.Utf8))
return df.with_columns(
pl.col(entity_col).map_dict(string_cache, return_dtype=pl.Int32)
pl.col(entity_col).replace(string_cache, return_dtype=pl.Int32, default=None)
)


def _reset_string_cache(
df: pl.DataFrame, inv_string_cache: Mapping[int, Union[int, str]], return_dtype
) -> pl.DataFrame:
return df.with_columns(
pl.col(df.columns[0]).map_dict(inv_string_cache, return_dtype=return_dtype)
pl.col(df.columns[0]).replace(
inv_string_cache, return_dtype=return_dtype, default=None
)
)


Expand Down

0 comments on commit 0d41c1f

Please sign in to comment.