pyspark - 选择客户在 pyspark 的 2 个字典中购买商品

Question

我想过滤在 list_1 中购买了这两种商品的客户和list_2或者同时购买了两者的客户 S/M和L/XL来自一 list_1或者同时购买了两者的客户S/M和L/XL来自一 list_2

list_1 = [A1, A2, B1, B2, C1, C2, D1, D2]

list_2 = [E1, E2, F1, F2, G1, G2, H1, H2]

<表类=“s-表”> <标题> customer_id 项目 尺寸 <正文> 1 A1 短/月 2 G1 短/月 3 D1 短/月 1 E2 L/XL 2 H2 L/XL 9 D1 短/月 1 G1 短/月 9 H2 L/XL 2 H2 L/XL

我想提取下表

<表类=“s-表”> <标题> customer_id 项目 <正文> 1 A1 2 G1 1 E2 9 D1 1 G1 9 H2 2 H2

我可以找到在 list_1 中购买了这两种商品的客户和list_2

w = Window.partitionBy('customer_id').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

df = (df.select('*',
                F.col('item').isin(list_1).alias('list_1'),
                F.col('item').isin(list_2).alias('list_2'))
      .select('customer_id', 'item',
              # If there is any True for a customer, take the True value
              F.max('list_1').over(w).alias('list_1'),
              F.max('list_2').over(w).alias('list_2'))
      .filter(F.col('list_1') & F.col('list_2'))
      .select('customer_id', 'item'))

但在过滤同时购买了这两种产品的客户时遇到了困难 S/M和L/XL来自一 list_1或者同时购买了两者的客户S/M和L/XL来自一 list_2 。我考虑创建一本字典

dict_1 = {"A1" : "S/M", "A2" : "L/XL", "B1" : "S/M", "B2" : "L/XL", "C1" : "S/M", "C2" : "L/XL","D1" : "S/M", "D2" : "L/XL"}

dict_2 = {"E1" : "S/M", "E2" : "L/XL", "F1" : "S/M", "F2" : "L/XL", "G1" : "S/M", "G2" : "L/XL", "H1" : "S/M", "H2" : "L/XL"}

我该如何处理这个问题？

Answer 1

请告诉我您对此解决方案的看法？

import pyspark.sql.functions as f

output_df = (
    df
    .withColumn('list_1', f.array(*[f.lit(element) for element in list_1]))
    .withColumn('list_2', f.array(*[f.lit(element) for element in list_2]))
    .withColumn('which_list', f.when(f.array_contains(f.col('list_1'), f.col('item')), f.lit('list_1')).otherwise(f.lit('list_2')))
    .withColumn('size_list', f.concat_ws('|', f.col('size'), f.col('which_list')))
    .groupBy('customer_id')
    .agg(
        f.collect_list(f.col('which_list')).alias('which_lists')
        f.collect_list(f.col('size_list')).alias('size_lists')
    )
    .where(
        (
            f.array_contains(f.col('which_lists'), 'lists_1') &
            f.array_contains(f.col('which_lists'), 'lists_2')
        ) |
        (
            f.array_contains(f.col('size_lists'), 'S/M|list_1') &
            f.array_contains(f.col('size_lists'), 'L/XL|list_1')
        ) |
        (
            f.array_contains(f.col('size_lists'), 'S/M|list_2') &
            f.array_contains(f.col('size_lists'), 'L/XL|list_2')
        )
    )
)