def add_row_indices(df, training=True):
  """
  Add a row index column for faster data ingestion times with SystemML.

  Args:
    df: A Spark DataFrame in which each row contains the slide number,
      tumor score, molecular score, and the sample stretched out into a
      Vector.
    training: Boolean for training or testing datasets.

  Returns:
    The Spark DataFrame with a row index column called "__INDEX".
  """
  rdd = (df.rdd
           .zipWithIndex()
           .map(lambda r: (r[1] + 1, *r[0])))  # flatten & convert index to 1-based indexing
  if training:
    df = rdd.toDF(['__INDEX', 'slide_num', 'tumor_score', 'molecular_score', 'sample'])
    df = df.select(df["__INDEX"].astype("int"), df.slide_num.astype("int"),
                   df.tumor_score.astype("int"), df.molecular_score, df["sample"])
  else:  # testing data -- no labels
    df = rdd.toDF(["__INDEX", "slide_num", "sample"])
    df = df.select(df["__INDEX"].astype("int"), df.slide_num.astype("int"), df["sample"])
  return df