PySpark Использование collect_list для сбора массивов различной длины

Я пытаюсь использовать collect_list для сбора массивов (и поддержания порядка) из двух разных фреймов данных.

Test_Data и Train_Data имеют одинаковый формат.

from pyspark.sql import functions as F
from pyspark.sql import Window

w = Window.partitionBy('Group').orderBy('date')

# Train_Data has 4 data points
# Test_Data has 7 data points
# desired target array:         [1, 1, 2, 3]
# desired MarchMadInd array:    [0, 0, 0, 1, 0, 0, 1]

sorted_list_diff_array_lens = train_data.withColumn('target', 
F.collect_list('target').over(w)
                                  )\
test_data.withColumn('MarchMadInd', F.collect_list('MarchMadInd').over(w))\
   .groupBy('Group')\
   .agg(F.max('target').alias('target'), 
    F.max('MarchMadInd').alias('MarchMadInd')
)

Я понимаю, что синтаксис неверен с "test_data.withColumn", но я хочу выбрать массив для MarchMadInd из test_date, но массив для target из train_data. Желаемый результат будет выглядеть следующим образом:

{"target":[1, 1, 2, 3], "MarchMadInd":[0, 0, 0, 1, 0, 0, 1]}

Контекст: это для модели временных рядов DeepAR (с использованием AWS), которая требует, чтобы динамические функции включали период прогнозирования, но целью должны быть исторические данные.