Pyspark - обновляет фрейм данных внутри цикла

Question

Я написал код, который обновляет фрейм данных внутри цикла for, но получаю странные результаты:

def _simulate_walks(self):

    # sample starting nodes
    aprox_sample_rate = 1.0 * self._num_of_walks / self._vertices.count()
    starting_nodes = self._vertices.sample(True, aprox_sample_rate)
    starting_nodes.show()

    # iterate over walks
    alias_draw_udf = F.udf(Node2Vec._alias_draw, T.StringType())
    single_list_udf = F.udf(lambda e: [e], T.ArrayType(T.StringType()))
    append_list_udf = F.udf(lambda l,e: l+[e], T.ArrayType(T.StringType()))
    for i in range(self._walk_length):
        if i == 0:
            chosen_path = starting_nodes.join(self._nodes_alias, F.col(self._src_col) == F.col('NODE'))\
                .withColumn('CHOSEN_NODE', alias_draw_udf('ALIAS'))\
                .withColumn('PATH', single_list_udf('NODE'))\
                .selectExpr('NODE as LAST_NODE', 'CHOSEN_NODE', 'PATH').persist()
            #chosen_path.show()
        else:
            chosen_path = chosen_path.join(self._edges_alias, (F.col('LAST_NODE') == F.col(self._src_col)) &
                                                (F.col('CHOSEN_NODE') == F.col(self._dst_col)))\
                .withColumn('NEW_CHOSEN_NODE', alias_draw_udf('ALIAS'))
            #chosen_path.show()
            chosen_path = chosen_path\
                .selectExpr('CHOSEN_NODE as LAST_NODE', 'NEW_CHOSEN_NODE as CHOSEN_NODE', 'PATH')
            #chosen_path.show()


        chosen_path = chosen_path.withColumn('NEW_PATH', append_list_udf('PATH', 'CHOSEN_NODE'))\
                .selectExpr('LAST_NODE', 'CHOSEN_NODE', 'NEW_PATH as PATH')
        chosen_path.show(5, False)

Однако, когда я добавляю постоянную команду внутри цикла:

chosen_path = chosen_path.withColumn('NEW_PATH', append_list_udf('PATH', 'CHOSEN_NODE'))\
                .selectExpr('LAST_NODE', 'CHOSEN_NODE', 'NEW_PATH as PATH').persist()

Код работает без нареканий.

Мне известно о том, что в искре код лениво оценивается до тех пор, пока действие не будет выполнено, но я бы не подумал, что это приведет к неожиданным результатам.

Само собой разумеется, что сохранение на каждой итерации не является хорошим решением из-за высокого потребления памяти, и мне интересно, как лучше всего решить эту проблему (может быть, не сохранять всю память и сразу после сохранения нового кадра данных?).

Спасибо.

Answer 1

Почему вы не используете functools.reduce ()? Пример:

from functools import reduce
def join_all_dataframes(dfs: List[DataFrame], on_columns: List[str]) -> DataFrame
   return reduce(lambda x, y: x.join(y, on_columns, 'outer'), 
                 dfs).dropDuplicates()