在Python类中注册Spark SQL用户定义函数

Question

我正在通过Python API使用Spark进行一些数据处理。下面是我的工作之类的简化一下：

class data_processor(object): 
    def __init__(self,filepath): 
     self.config = Config() # this loads some config options from file 
     self.type_conversions = {int:IntegerType,str:StringType} 
     self.load_data(filepath) 
     self.format_age() 

    def load_data(self,filepath,delim='\x01'): 
     cols = [...] # list of column names 
     types = [int, str, str, ... ] # list of column types 
     user_data = sc.textFile(filepath,use_unicode=False).map(lambda row: [types[i](val) for i,val in enumerate(row.strip().split(delim))]) 
     fields = StructType([StructField(field_name,self.type_conversions[field_type]()) for field_name,field_type in zip(cols,types)]) 
     self.user_data = user_data.toDF(fields) 
     self.user_data.registerTempTable('data') 

    def format_age(self): 
     age_range = self.config.age_range # tuple of (age_min, age_max) 
     age_bins = self.config.age_bins # list of bin boundaries 
     def _format_age(age): 
      if age<age_range[0] or age>age_range[1]: 
       return None 
      else: 
       return np.digitize([age],age_bins)[0] 
     sqlContext.udf.register('format_age', lambda x: _format_age(x), IntegerType()) 

现在，如果我实例化data=data_processor(filepath)类的，我可以做查询的数据帧就好了。例如，这个作品：

sqlContext.sql("select * from data limit 10").take(1) 

但我很明显没有正确设置udf。如果我尝试，例如，

sqlContext.sql("select age, format_age(age) from data limit 10").take(1) 

我得到一个错误：

Py4JJavaError: An error occurred while calling z:org.apache.spark.api.python.PythonRDD.collectAndServe. 

（用长堆栈跟踪，典型的星火，时间太长，包括在这里）。

那么，我究竟做错了什么？在像这样的方法中定义UDF的方法是什么（最好是作为类方法）。我知道Spark不喜欢传递类对象，因此嵌套结构为format_age（受this question启发）。

想法？

Answer 1

答案很简单。您不能使用NumPy数据类型作为Spark SQL中标准Python类型的直接替换。当您声明返回类型为IntegerType时，返回类型np.digitize是numpy.int64而非int。

所有你需要做的就是投从_format_age返回的值：

def _format_age(age): 
    ... 
    return int(np.digitize(...))