Apache Spark 如何定义一个泛型函数来从一个给定的列名派生多个新列?

jpfvwuh4  于 2022-11-16  发布在  Apache
关注(0)|答案(1)|浏览(107)

我一直在尝试通过一个python函数从一个给定的列名派生出两个列。下面是代码片段:

from pyspark.sql.functions import substring

def deriveCol(source_col_name, col1, col2):
    df.select(source_col_name, substring(source_col_name, 1, 4).alias(col1), substring(source_col_name, 5, 2).alias(col2))
    for i in col2:
        if i >= "01" and i <= "03":
            print("First quarter")
        elif i >= "04" and i <= "06":
            print("second quarter")
        elif i >= "06" and i <= "09":
            print("Third quarter")
        else:
            print("Fourth quarter")
    return df.select([col1, col2]).show(10, truncate=True)
 
t = deriveCol("Report", "year", "month")

请参见以下输出

Fourth quarter
Fourth quarter
+----+-----+
|year|month|
+----+-----+
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
|2022|   01|
+----+-----+
only showing top 10 rows

“Report' =必须从中导出数据的源列。
“年”/“月”=派生列。
我有一个数据库表的列,我必须将其拆分为3个部分:“年”、“月”、“季度”。
“year”和“month”很简单。但是找到“quarter”不起作用。当执行代码时,它直接进入else并打印“fourth quarter”。但是我可以说,我在上面substring()函数的帮助下导出的“month”也有来自给定年份所有月份的日期。
注意:我需要将“季度”的输出保存在单独的列中。

r6l8ljro

r6l8ljro1#

作为第一步,我已经从 string 中创建了一个真正的 date 列。有了 date 列,我们可以应用方法yearmonthquarter

def derive_col(source_col_name, col1, col2, col3):
    date_col = F.to_date(source_col_name, 'yyyyMMdd')
    df = df.select(
        *[c for c in df.columns if c != source_col_name],
        F.year(date_col).alias(col1),
        F.month(date_col).alias(col2),
        F.quarter(date_col).alias(col3)
    )
    return df

完整测试,使用多个df参数,因为没有它我无法使用函数:

from pyspark.sql import functions as F

df = spark.createDataFrame([('20220101',), ('20220731',)], ['Report'])

def derive_col(df, source_col_name, col1, col2, col3):
    date_col = F.to_date(source_col_name, 'yyyyMMdd')
    df = df.select(
        *[c for c in df.columns if c != source_col_name],
        F.year(date_col).alias(col1),
        F.month(date_col).alias(col2),
        F.quarter(date_col).alias(col3)
    )
    return df

df = derive_col(df, 'Report', 'year', 'month', 'quarter')

df.show()
# +----+-----+-------+
# |year|month|quarter|
# +----+-----+-------+
# |2022|    1|      1|
# |2022|    7|      3|
# +----+-----+-------+

相关问题