QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-21 00:54:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: drop the first level of MultiIndex (#19716)

Browse Source
This commit is contained in:
Yongjie Zhao
2022-04-14 23:40:38 +08:00
committed by GitHub
parent 2b17ce2bd0
commit 9425dd2cac
5 changed files with 140 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
tests/unit_tests/pandas_postprocessing/test_flatten.py
View File

@@ -18,6 +18,7 @@ import pandas as pd
from superset.utils import pandas_postprocessing as pp
from superset.utils.pandas_postprocessing.utils import FLAT_COLUMN_SEPARATOR
from tests.unit_tests.fixtures.dataframes import timeseries_df
def test_flat_should_not_change():
@@ -73,3 +74,85 @@ def test_flat_should_flat_multiple_index():
}
)
)
def test_flat_should_drop_index_level():
index = pd.to_datetime(["2021-01-01", "2021-01-02", "2021-01-03"])
index.name = "__timestamp"
columns = pd.MultiIndex.from_arrays(
[["a"] * 3, ["b"] * 3, ["c", "d", "e"], ["ff", "ii", "gg"]],
names=["level1", "level2", "level3", "level4"],
)
df = pd.DataFrame(index=index, columns=columns, data=1)
# drop level by index
assert pp.flatten(df.copy(), drop_levels=(0, 1,)).equals(
pd.DataFrame(
{
"__timestamp": index,
FLAT_COLUMN_SEPARATOR.join(["c", "ff"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["d", "ii"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["e", "gg"]): [1, 1, 1],
}
)
)
# drop level by name
assert pp.flatten(df.copy(), drop_levels=("level1", "level2")).equals(
pd.DataFrame(
{
"__timestamp": index,
FLAT_COLUMN_SEPARATOR.join(["c", "ff"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["d", "ii"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["e", "gg"]): [1, 1, 1],
}
)
)
# only leave 1 level
assert pp.flatten(df.copy(), drop_levels=(0, 1, 2)).equals(
pd.DataFrame(
{
"__timestamp": index,
FLAT_COLUMN_SEPARATOR.join(["ff"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["ii"]): [1, 1, 1],
FLAT_COLUMN_SEPARATOR.join(["gg"]): [1, 1, 1],
}
)
)
def test_flat_should_not_droplevel():
assert pp.flatten(timeseries_df, drop_levels=(0,)).equals(
pd.DataFrame(
{
"index": pd.to_datetime(
["2019-01-01", "2019-01-02", "2019-01-05", "2019-01-07"]
),
"label": ["x", "y", "z", "q"],
"y": [1.0, 2.0, 3.0, 4.0],
}
)
)
def test_flat_integer_column_name():
index = pd.to_datetime(["2021-01-01", "2021-01-02", "2021-01-03"])
index.name = "__timestamp"
columns = pd.MultiIndex.from_arrays(
[["a"] * 3, [100, 200, 300]],
names=["level1", "level2"],
)
df = pd.DataFrame(index=index, columns=columns, data=1)
assert pp.flatten(df, drop_levels=(0,)).equals(
pd.DataFrame(
{
"__timestamp": pd.to_datetime(
["2021-01-01", "2021-01-02", "2021-01-03"]
),
"100": [1, 1, 1],
"200": [1, 1, 1],
"300": [1, 1, 1],
}
)
)