Skip to content

Commit

Permalink
Remove unused levels from categorical index
Browse files Browse the repository at this point in the history
If `unique_id` is categorical, and some levels were removed, then
`value_counts()` in `_grouped_array_from_df()` still returns these
levels in the output with their counts being 0.

`_grouped_array_from_df()`, however, does not expect that some counts
could be zero, therefore, it makes sense to drop unused levels from a
categorical index.
  • Loading branch information
nickto committed Apr 21, 2023
1 parent 3aac9c5 commit 9bd568b
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 0 deletions.
5 changes: 5 additions & 0 deletions nbs/core.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -996,6 +996,11 @@
" self._validate_df(df)\n",
" if df.index.name != 'unique_id':\n",
" df = df.set_index('unique_id')\n",
" if isinstance(df.index, pd.CategoricalIndex):\n",
" # If `unique_id` is categorical and some levels are not present\n",
" # in `df`, then `.count_values()` in `unique_id` does not\n",
" # produce expected output, so drop these unused levels.\n",
" df.index = df.index.remove_unused_categories()\n",
" df = _parse_ds_type(df)\n",
" self.ga, self.uids, self.last_dates, self.ds = _grouped_array_from_df(df, sort_df)\n",
" self.n_jobs = _get_n_jobs(len(self.ga), self.n_jobs)\n",
Expand Down
5 changes: 5 additions & 0 deletions statsforecast/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,6 +547,11 @@ def _prepare_fit(self, df, sort_df):
self._validate_df(df)
if df.index.name != "unique_id":
df = df.set_index("unique_id")
if isinstance(df.index, pd.CategoricalIndex):
# If `unique_id` is categorical and some levels are not present
# in `df`, then `.count_values()` in `unique_id` does not
# produce expected output, so drop these unused levels.
df.index = df.index.remove_unused_categories()
df = _parse_ds_type(df)
self.ga, self.uids, self.last_dates, self.ds = _grouped_array_from_df(
df, sort_df
Expand Down

0 comments on commit 9bd568b

Please sign in to comment.