pyspark数据帧上的自定义功能

时间:2017-12-01 16:39:55

标签: python pyspark

我尝试在pyspark数据帧中的行上应用自定义函数。 此函数采用行和其他2个相同维度的向量。它为第二个向量中的行输出每个匹配值的第三个向量的值之和。

    int row = tab_reta_in.getSelectedRow();
    int col = tab_reta_in.getSelectedColumn();

    Data_feild = (String) tab_reta_in.getValueAt(row, col);
    int Def_r = row;


    if (col == 0) {

        try {
            ResultSet rs = j.GetData("select * from invoice_retails where name like '" + Data_feild + "%' ");

            while (rs.next()) {

                String col0 = rs.getString("medi_id");
                String col1 = rs.getString("name");
                String col2 = rs.getString("prize");

                Object[] rows = {col0, col1, col2};
                dt.addRow(rows);
            }

        } catch (Exception ex) {
            System.out.println(ex);
        }

功能:

import pandas as pd
import numpy as np

我想通过熊猫实现的目标很简单:

def V_sum(row,b,c):
    return float(np.sum(c[row==b]))

我想在pyspark中执行相同的操作。

pd_df = pd.DataFrame([[0,1,0,0],[1,1,0,0],[0,0,1,0],[1,0,1,1],[1,1,0,0]], columns=['t1', 't2', 't3', 't4'])
   t1  t2  t3  t4
0   0   1   0   0
1   1   1   0   0
2   0   0   1   0
3   1   0   1   1
4   1   1   0   0

B = np.array([1,0,1,0])
V = np.array([5,1,2,4])

pd_df.apply(lambda x: V_sum(x, B, V), axis=1)
0    4.0
1    9.0
2    7.0
3    8.0
4    9.0
dtype: int64

我考虑过使用udf,但我无法使用它

from pyspark import SparkConf, SparkContext, SQLContext
sc = SparkContext("local")
sqlContext = SQLContext(sc)

spk_df = sqlContext.createDataFrame([[0,1,0,0],[1,1,0,0],[0,0,1,0],[1,0,1,1],[1,1,0,0]], ['t1', 't2', 't3', 't4'])
spk_df.show()
+---+---+---+---+
| t1| t2| t3| t4|
+---+---+---+---+
|  0|  1|  0|  0|
|  1|  1|  0|  0|
|  0|  0|  1|  0|
|  1|  0|  1|  1|
|  1|  1|  0|  0|
+---+---+---+---+

显然,我做错了,因为它产生了:

from pyspark.sql.types import FloatType
import pyspark.sql.functions as F

V_sum_udf = F.udf(V_sum, FloatType()) 
spk_df.select(V_sum_udf(F.array(*(F.col(x) for x in spk_df.columns))).alias("results")).show()

1 个答案:

答案 0 :(得分:0)

如果您希望在函数内部使用非列数据以及列数据来计算新列,那么如here所述的UDF +闭包+ withColumn是一个很好的起点。

B = [2,0,1,0] 
V = [5,1,2,4]

v_sum_udf = F.udf(lambda row: V_sum(row, B, V), FloatType())
spk_df.withColumn("results", v_sum_udf(F.array(*(F.col(x) for x in spk_df.columns))))