code: def create_hier_dataset(
      spark: pyspark.sql.SparkSession,
      df: pyspark.sql.DataFrame,
      time_col: str = None,
      index_cols: List[str] = None,
      agg_dict: Dict = None,
  ) -> Tuple[pyspark.sql.DataFrame, np.ndarray]:
      """
      Aggregates the time series in the dataset & appends them to the original dataset.
  
      :param spark: The current SparkSession.
      :param df: A pyspark dataframe containing all the data. The dataframe must have a column `TSID_COL_NAME`
          indexing the time series in the dataset (this column is automatically added by `read_dataset`).
      :param time_col: The name of the column which specifies timestamp. If ``None`` is provided, it is assumed to be the
          first column which is not an index column or pre-specified data column.
      :param index_cols: The columns used to index the various time series in the dataset. If ``None`` is provided, we
          assume the entire dataset is just a single time series. These columns define the levels of the hierarchy.
          For example, if each time series represents sales and we have ``index_cols = ["store", "item"]``, we will
          first aggregate sales for all items sold at a particular store; then we will aggregate sales for all items at
          all stores.
      :param agg_dict: A dictionary used to specify how different data columns should be aggregated. If a data column
          is not in the dict, we aggregate using sum by default.
  
      :return: The dataset with additional time series corresponding to each level of the hierarchy, as well as a
          matrix specifying how the hierarchy is constructed.
      """
      # Determine which columns are index vs. data columns
      index_cols = [] if index_cols is None else index_cols
      index_cols = [c for c in index_cols if c != TSID_COL_NAME]
      extended_index_cols = index_cols + [TSID_COL_NAME]
      if time_col is None:
          non_index_cols = [
              c
              for c in df.schema.fieldNames()
              if c not in extended_index_cols
          ]
          time_col = non_index_cols[0]
          data_cols = non_index_cols[1:]
      else:
          data_cols = [
              c
              for c in df.schema.fieldNames()
              if c not in extended_index_cols + [time_col]
          ]
  
      # Create a pandas index for all the time series
      ts_index = (
          df.groupBy(extended_index_cols)
          .count()
          .drop("count")
          .toPandas()
      )
      ts_index = ts_index.set_index(index_cols).sort_index()
      index_schema = StructType(
          [df.schema[c] for c in extended_index_cols]
      )
      n = len(ts_index)
  
      # Add all higher levels of the hierarchy
      full_df = df
      hier_vecs = []
  
      # Compose the aggregation portions of the SQL select statements below
      df.createOrReplaceTempView("df")
      agg_dict = {} if agg_dict is None else agg_dict
      data_col_sql = [
          f"{agg_dict.get(c, 'sum').upper()}(`{c}`) AS `{c}`"
          for c in data_cols
      ]
  
      for k in range(len(index_cols)):
          # Aggregate values of data columns over the last k+1 index column values.
          gb_cols = index_cols[: -(k + 1)]
          gb_col_sql = [f"`{c}`" for c in [time_col] + gb_cols]
          agg = spark.sql(
              f"SELECT {','.join(gb_col_sql + data_col_sql)} FROM df GROUP BY {','.join(gb_col_sql)};"
          )
  
          # Add back dummy NA values for the index columns we aggregated over, add a time series ID column,
          # concatenate the aggregated time series to the full dataframe, and compute the hierarchy vector.
  
          # For the top level of the hierarchy, this is easy as we just sum everything
          if len(gb_cols) == 0:
              dummy = [
                  ["__aggregated__"] * len(index_cols)
                  + [n + len(hier_vecs)]
              ]
              full_df = full_df.unionByName(
                  agg.join(
                      spark.createDataFrame(dummy, schema=index_schema)
                  )
              )
              hier_vecs.append(np.ones(n))
              continue
  
          # For lower levels of the hierarchy, we determine the membership of each grouping to create
          # the appropriate dummy entries and hierarchy vectors.
          dummy = []
          for i, (group, group_idxs) in enumerate(
              ts_index.groupby(gb_cols).groups.items()
          ):
              group = [group] if len(gb_cols) == 1 else list(group)
              locs = [ts_index.index.get_loc(j) for j in group_idxs]
              dummy.append(
                  group
                  + ["__aggregated__"] * (k + 1)
                  + [n + len(hier_vecs)]
              )
              x = np.zeros(n)
              x[locs] = 1
              hier_vecs.append(x)
          dummy = spark.createDataFrame(dummy, schema=index_schema)
          full_df = full_df.unionByName(agg.join(dummy, on=gb_cols))
  
      # Create the full hierarchy matrix, and return it along with the updated dataframe
      hier_matrix = np.concatenate([np.eye(n), np.stack(hier_vecs)])
      return full_df, hier_matrix
  
imports: from typing import Dict, List, Tuple
  import numpy
  import pandas
label_desc: Use image node_id 3 for calling a global function with example usage: create_hier_dataset(spark, df, time_col, index_cols, agg_dict) and returns: full_df,  hier_matrix