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

Question

Допустим, у меня есть две группы: '1' и '2'.

ls = [
    ['1', 2],
    ['2', 7],
    ['1', 3],
    ['2',-6],
    ['1', 3],
    ['1', 5],
    ['1', 4],
    ['2', 7]
]
df = spark.createDataFrame(pd.DataFrame(ls, columns=['col1', 'col2']))
df.show()

+----+-----+
|col1| col2|
+----+-----+
|   1|    2|
|   2|    7|
|   1|    3|
|   2|   -6|
|   1|    3|
|   1|    5|
|   1|    4|
|   2|    7|
+----+-----+

Я рассчитал QuantileDiscretizer () для каждой группы, используя:

discretizer = QuantileDiscretizer(numBuckets=3, inputCol="col2", outputCol="result")
bucketizer_1 = discretizer.fit(df.filter(col('col1')=='1'))
bucketizer_2 = discretizer.fit(df.filter(col('col1')=='2'))

Эти сгустители хранятся в списке, поэтому я могу назвать их с помощью индексации (сгустители [0], сгустители [1])

Теперь я пытаюсь преобразовать данные в столбце col2 в квантили, используя бакетизаторы. Другими словами, если col1 == '1', то сглаживатели [0] .transform (data) -> сохранят результат в новом столбце кадра данных. Тот же процесс для группы «2», но с добавками [1].

(Давайте не будем обращать внимание на причину, по которой я не уместился, и преобразовать данные за один раз)

Я изо всех сил пытаюсь написать код, чтобы получить то, что я хочу, так как transform () использует столбец в качестве входных данных.

Может кто-нибудь помочь?

Answer 1

Одно решение, которое я могу придумать, - это использование union :

bucketizers[0].transform(df.filter((F.col('col1')=='1'))).union(bucketizers[1].transform(df.filter(F.col('col1')=='2'))).show()