feat(advanced analysis): support MultiIndex column in post processing stage (#19116)

superset/utils/pandas_postprocessing/cum.py

@@ -14,27 +14,21 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Any, Dict, Optional
+from typing import Dict
 
 from flask_babel import gettext as _
 from pandas import DataFrame
 
-from superset.exceptions import QueryObjectValidationError
+from superset.exceptions import InvalidPostProcessingError
 from superset.utils.pandas_postprocessing.utils import (
     _append_columns,
-    _flatten_column_after_pivot,
     ALLOWLIST_CUMULATIVE_FUNCTIONS,
     validate_column_args,
 )
 
 
 @validate_column_args("columns")
-def cum(
-    df: DataFrame,
-    operator: str,
-    columns: Optional[Dict[str, str]] = None,
-    is_pivot_df: bool = False,
-) -> DataFrame:
+def cum(df: DataFrame, operator: str, columns: Dict[str, str],) -> DataFrame:
     """
     Calculate cumulative sum/product/min/max for select columns.
 
@@ -45,29 +39,16 @@ def cum(
            `y2` based on cumulative values calculated from `y`, leaving the original
            column `y` unchanged.
     :param operator: cumulative operator, e.g. `sum`, `prod`, `min`, `max`
-    :param is_pivot_df: Dataframe is pivoted or not
     :return: DataFrame with cumulated columns
     """
     columns = columns or {}
-    if is_pivot_df:
-        df_cum = df
-    else:
-        df_cum = df[columns.keys()]
+    df_cum = df.loc[:, columns.keys()]
     operation = "cum" + operator
     if operation not in ALLOWLIST_CUMULATIVE_FUNCTIONS or not hasattr(
         df_cum, operation
     ):
-        raise QueryObjectValidationError(
+        raise InvalidPostProcessingError(
             _("Invalid cumulative operator: %(operator)s", operator=operator)
         )
-    if is_pivot_df:
-        df_cum = getattr(df_cum, operation)()
-        agg_in_pivot_df = df.columns.get_level_values(0).drop_duplicates().to_list()
-        agg: Dict[str, Dict[str, Any]] = {col: {} for col in agg_in_pivot_df}
-        df_cum.columns = [
-            _flatten_column_after_pivot(col, agg) for col in df_cum.columns
-        ]
-        df_cum.reset_index(level=0, inplace=True)
-    else:
-        df_cum = _append_columns(df, getattr(df_cum, operation)(), columns)
+    df_cum = _append_columns(df, getattr(df_cum, operation)(), columns)
     return df_cum