как использовать значения последовательных строк сгруппированных данных для определения значения текущей строки с помощью pyspark

Question

Из набора данных ниже я хочу изменить значение столбца won_offer на 1 или 0. Проблема в том, что мне нужны последовательные строки комбинации кода клиента, чтобы определить значение этого столбца.

Если одна из следующих строк в течение 30 дней с даты текущей строки содержит order и цена ниже, чем цена текущей строки, 0 столбца won_offer этой строки может стать 1.

образец набора данных:

analysis = sqlContext.createDataFrame(
    [
        ('customer1', 'code1', 'date', 'order', 1.7, 0, 1),
        ('customer1', 'code2', 'date', 'offer', 1.5, 0, 2),
        ('customer1', 'code2', 'date', 'offer', 2.0, 0, 2),
        ('customer2', 'code1', 'date', 'offer', 1.2, 0,4),
        ('customer2', 'code1', 'date', 'order', 1.1, 0,4),
        ('customer2', 'code1', 'date', 'order', 2.0, 0,4),
        ('customer2', 'code1', 'date', 'offer', 1.2, 0,4)
    ],
    ('customer', 'code', 'order_date', 'type', 'price', 'final_offer', 'counter')
)

Я пробовал что-то подобное, но это не работает, так как я не знаю, как передать несколько строк в мой udf:

w = \
    Window.partitionBy('customer','code').orderBy('orderoffer_date')

@F.udf(returnType=IntegerType())
def logic_udf(counter, curr_date, next_dates, current_type, next_types, curr_price, next_prices) :
    for i in range(len(counter)):
        if (next_dates[i] < curr_date+30):
            if (next_types[i] == 'order') & (next_prices[i] < curr_price ):
                return 1
            else:
                return 0
        else:
            return 0

analysis = analysis.withColumn('won_offer', 
               logic(analysis.counter, analysis.order_date,lead(analysis.order_date, 
               analysis.n).over(w), analysis.type,lead(analysis.type, 
               analysis.n).over(w), analysis.price, lead(analysis.price, 
               analysis.n).over(w)))

желаемый вывод:

desired_result = sqlCtx.createDataFrame(
    [
        ('customer1', 'code1', 'date', 'order', 1.7, 0, 1),
        ('customer1', 'code2', 'date', 'offer', 1.5, 0, 2),
        ('customer1', 'code2', 'date', 'offer', 2.0, 0, 2),
        ('customer2', 'code1', 'date', 'offer', 1.2, 1, 4),
        ('customer2', 'code1', 'date', 'order', 1.1, 1, 4),
        ('customer2', 'code1', 'date', 'order', 1.0, 0, 4),
        ('customer2', 'code1', 'date', 'offer', 1.2, 0, 4)
    ],
    ('customer', 'code', 'order_date', 'type', 'price', 'final_offer', 'counter')
)

Я понимаю, что мой вопрос довольно сложный. Если бы кто-то мог просто сказать мне, как я мог бы передать несколько строк сгруппированных данных в udf , мне бы уже очень помогли.

Короче говоря: основная цель - определить значение столбца в строке, просматривая несколько столбцов в следующих строках (и все же в их определенной группе).

Заранее спасибо! Charles

Answer 1

Вы можете использовать окно с функциями sql для замены logic_udf. Поскольку вы используете только первую строку после текущей строки, вы можете добавить первую строку после текущей строки в текущую строку.

from pyspark.sql import functions as F

analysis \
    .withColumn('next_order_date', F.first('order_date').over(w)) \
    .withColumn('next_type', F.first('type').over(w)) \
    .withColumn('next_price', F.first('price').over(w)) \
    .withColumn('won_offer', F.when(condition, 1).otherwise(0))