前填充多列可重用功能代码

时间:2019-05-10 16:08:51

标签: apache-spark pyspark pyspark-sql

我正在尝试将基于先前堆栈溢出后的正向填充估算过程转换为可重用函数(带有def(...)的东西),因此我可以将其应用于多个列,而不是每个列都有代码段柱。使用参数创建可重用的函数一直是我的挑战。

谢谢!

帖子=> Forward fill missing values in Spark/Python

代码示例代码段

# sample data
df = spark.createDataFrame([('2019-05-10 7:30:05', '10', '0.5'),\
                            ('2019-05-10 7:30:10', 'UNKNOWN', '0.24'),\
                            ('2019-05-10 7:30:15', '6', 'UNKNOWN'),\
                            ('2019-05-10 7:30:20', '7', 'UNKNOWN'),\
                            ('2019-05-10 7:30:25', '10', '1.1'),\
                            ('2019-05-10 7:30:30', 'UNKNOWN', '1.1'),\
                            ('2019-05-10 7:30:35', 'UNKNOWN', 'UNKNOWN'),\
                            ('2019-05-10 7:30:49', '50', 'UNKNOWN')], ["date", "v1", "v2"])

df = df.withColumn("date", F.col("date").cast("timestamp"))

# schema
root
 |-- date: timestamp (nullable = true)
 |-- v1: string (nullable = true)
 |-- v2: string (nullable = true)

# imputer process / all cols that need filled are strings
def stringReplaceFunc(x, y):
    '''
    this function replaces column values:
    ex: replace 'UNKNOWN' reading with nulls for forward filling function
    : x => source col
    : y => replace value
    '''
    return F.when(x != y, x).otherwise(F.lit(None)) # replace with NULL

# this windows function triggers forward filling for null values created from StringReplaceFunc
window = Window\
.partitionBy(F.month("date"))\
.orderBy('date')\
.rowsBetween(-sys.maxsize, 0)

# here is where I am trying to make a function so I don't have to code each col that needs filled individually
df = df\
.withColumn("v1", stringReplaceFunc(F.col("v1"), "UNKNOWN"))

fill_v1 = F.last(df['v1'], ignorenulls=True).over(window)
df = df.withColumn('v1',  fill_v1)

df = df\
.withColumn("v2", stringReplaceFunc(F.col("v2"), "UNKNOWN"))

fill_v1 = F.last(df['v2'], ignorenulls=True).over(window)
df = df.withColumn('v2',  fill_v1)

# imputing results of the output needed
df.show()

+-------------------+---+----+
|               date| v1|  v2|
+-------------------+---+----+
|2019-05-10 07:30:05| 10| 0.5|
|2019-05-10 07:30:10| 10|0.24|
|2019-05-10 07:30:15|  6|0.24|
|2019-05-10 07:30:20|  7|0.24|
|2019-05-10 07:30:25| 10| 1.1|
|2019-05-10 07:30:30| 10| 1.1|
|2019-05-10 07:30:35| 10| 1.1|
|2019-05-10 07:30:49| 50| 1.1|
+-------------------+---+----+

2 个答案:

答案 0 :(得分:1)

我不是100%正确地理解了这个问题,但这是将您提到的代码封装到python函数中的一种方式:

def forward_fill(df, col_name):
    df = df.withColumn(col_name, stringReplaceFunc(F.col(col_name), "UNKNOWN"))

    last_func = F.last(df[col_name], ignorenulls=True).over(window)
    df = df.withColumn(col_name,  last_func)
    return df

然后您可以将其称为:df = forward_fill(df, 'v1')

答案 1 :(得分:1)

这是可行的解决方案

def stringReplaceFunc(x, y):
    return F.when(x != y, x).otherwise(F.lit(None)) # replace with NULL

def forwardFillImputer(df, cols=[], partitioner="date", value="UNKNOWN"):
  for i in cols:
    window = Window\
    .partitionBy(F.month(partitioner))\
    .orderBy(partitioner)\
    .rowsBetween(-sys.maxsize, 0)
    df = df\
    .withColumn(i, stringReplacer(F.col(i), value))
    fill = F.last(df[i], ignorenulls=True).over(window)
    df = df.withColumn(i,  fill)
  return df
df = forwardFillImputer(df, cols=[i for i in df.columns])