pyspark.pandas .apply() changing rows ordering

When using the apply function to apply a function to a DataFrame column, it ends up mixing the column's rows ordering.

A command like this:

def example_func(df_col):
  return df_col ** 2 

df['col_to_apply_function'] = df.apply(lambda row: example_func(row['col_to_apply_function']), axis=1) 

A workaround is to assign the results to a new column instead of the same one, but if the old column is dropped, the same error is produced.

Setting one column as index also didn't work.