Skip to content

Commit

Permalink
fix: add check for length in TSDataset.to_hierarchical_dataset, add t…
Browse files Browse the repository at this point in the history
…est for it
  • Loading branch information
d.a.bunin committed Apr 18, 2023
1 parent 77604f9 commit 468774b
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 1 deletion.
10 changes: 9 additions & 1 deletion etna/datasets/tsdataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -812,10 +812,18 @@ def to_hierarchical_dataset(
-------
:
Dataframe in wide format and optionally hierarchical structure
Raises
------
ValueError
If ``level_columns`` is empty
"""
if len(level_columns) == 0:
raise ValueError("Value of level_columns shouldn't be empty!")

df_copy = df.copy(deep=True)
df_copy["segment"] = df_copy[level_columns].astype("string").agg(sep.join, axis=1)
if not keep_level_columns and len(level_columns) > 0:
if not keep_level_columns:
df_copy.drop(columns=level_columns, inplace=True)
df_copy = TSDataset.to_dataset(df_copy)

Expand Down
6 changes: 6 additions & 0 deletions tests/test_datasets/test_hierarchical_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -307,6 +307,12 @@ def test_level_names_without_hierarchical_structure(market_level_df):
assert ts_level_names is None


def test_to_hierarchical_dataset_fails_empty_level_columns(product_level_df_long):
df = product_level_df_long
with pytest.raises(ValueError, match="Value of level_columns shouldn't be empty"):
_ = TSDataset.to_hierarchical_dataset(df=df, level_columns=[])


def test_to_hierarchical_dataset_not_change_input_df(product_level_df_long):
df = product_level_df_long
df_before = df.copy()
Expand Down

0 comments on commit 468774b

Please sign in to comment.