Pandas, объединение двух DataFrame

Question

Есть стандартный датасет Титаника: train и test

Хочу подготовить данные для моделей, модифицируя их, и, соответственно, есть желание делать преобразования сразу в test и train

full_data = [train, test]

Далее запускаю циклы, например:

for dataset in full_data:
    dataset['*col_name*'] = dataset['*col_name*'].apply(function) 

Они работают и если я вызываю train или test, то я получаю DF после необходимых мне преобразований

Проблема возникла вот в этом коде:

for dataset in full_data:
    polym = pd.DataFrame(poly.fit_transform(dataset[['Age', 'Fare']]))
    dataset = pd.concat([dataset, polym], axis=1)

В результате, если я вызову объект train, то он будет без изменений, а если вызову объект dataset, то он как раз будет содержать DF того вида, которого мне нужно (будут добавлены нелинейные преобразования некоторых признаков)

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

P.S. думал, что может проблемы с pd.concat(), но вроде бы ясно как он работает. Не могу понять почему преобразованный DF он записывает на новую переменную dataset, а не на train / test

Спасибо!