过滤pyspark数据框中的行并创建一个包含结果的新列

时间:2020-06-29 22:37:59

标签: sql pyspark user-defined-functions

因此,我试图确定周日在SF市区边界内发生的犯罪。我的想法是首先编写一个UDF来标记是否每项犯罪都在我确定为市区的区域内,如果在该区域内发生,则标记为“ 1”,否则为“ 0”。之后,我试图创建一个新列来存储这些结果。我尽力写了所有我能写的东西,但是由于某种原因它没有用。这是我写的代码:

from pyspark.sql.types import BooleanType
from pyspark.sql.functions import udf

def filter_dt(x,y):
  if (((x < -122.4213) & (x > -122.4313)) & ((y > 37.7540) & (y < 37.7740))):
    return '1'
  else:
    return '0'

schema = StructType([StructField("isDT", BooleanType(), False)])
filter_dt_boolean = udf(lambda row: filter_dt(row[0], row[1]), schema)

#First, pick out the crime cases that happens on Sunday BooleanType()
q3_sunday = spark.sql("SELECT * FROM sf_crime WHERE DayOfWeek='Sunday'")
#Then, we add a new column for us to filter out(identify) if the crime is in DT
q3_final = q3_result.withColumn("isDT", filter_dt(q3_sunday.select('X'),q3_sunday.select('Y')))

我得到的错误是:Picture for the error message

我的猜测是,我现在拥有的udf不支持将整个列作为要比较的输入,但是我不知道如何对其进行修复以使其起作用。请帮忙!谢谢!

2 个答案:

答案 0 :(得分:1)

尝试如下更改最后一行-

from pyspark.sql.functions import col
q3_final = q3_result.withColumn("isDT", filter_dt(col('X'),col('Y')))

答案 1 :(得分:0)

样本数据会有所帮助。现在,我假设您的数据如下所示:

+----+---+---+
|val1|  x|  y|
+----+---+---+
|  10|  7| 14|
|   5|  1|  4|
|   9|  8| 10|
|   2|  6| 90|
|   7|  2| 30|
|   3|  5| 11|
+----+---+---+

那么您就不需要udf,因为您可以使用when()函数进行评估

import pyspark.sql.functions as F
tst= sqlContext.createDataFrame([(10,7,14),(5,1,4),(9,8,10),(2,6,90),(7,2,30),(3,5,11)],schema=['val1','x','y'])
tst_res = tst.withColumn("isdt",F.when(((tst.x.between(4,10))&(tst.y.between(11,20))),1).otherwise(0))This will give the result
   tst_res.show()
+----+---+---+----+
|val1|  x|  y|isdt|
+----+---+---+----+
|  10|  7| 14|   1|
|   5|  1|  4|   0|
|   9|  8| 10|   0|
|   2|  6| 90|   0|
|   7|  2| 30|   0|
|   3|  5| 11|   1|
+----+---+---+----+

如果我的数据有误,但仍然需要将多个值传递给udf,则必须将其作为数组或结构传递。我更喜欢一个结构

from pyspark.sql.functions import udf
from pyspark.sql.types import *

@udf(IntegerType())
def check_data(row):
    if((row.x in range(4,5))&(row.y in range(1,20))):
        return(1)
    else:
        return(0)
tst_res1 = tst.withColumn("isdt",check_data(F.struct('x','y')))

结果将是相同的。但是,最好避免使用UDF并使用Spark内置函数,因为Spark催化剂无法理解udf内部的逻辑,也无法对其进行优化。

相关问题