使用udf的Pyspark进程数组列并返回另一个数组

时间:2019-02-06 17:37:56

标签: apache-spark pyspark

使用udf处理数组列并返回另一个数组

下面是我的输入内容:

docID带状疱疹 D1 [23,25,39,59] D2 [34、45、65]

我想通过处理带状疱疹数组列来生成一个称为哈希的新列: 例如,我要提取min和max(这只是示例,它表明我想要一个固定长度的数组列,而我实际上并不希望找到min或max)

docID带状疱疹 D1 [23,25,39,59] [23,59] D2 [34,45,65] [34,65]

我创建了如下的udf:

python manage.py test integration_tests/integration_*  --noinput --testrunner=lib.nodb_test_runner.NoDbTestRunner

但是它给出了以下错误:

def generate_minhash_signatures(shingles, coeffA, coeffB):
    signature = []
    minHashCode = nextPrime + 1
    maxHashCode = 0
    for shingleID in shingles:
        if shingleID < minHashCode:
            minHashCode = shingleID
        if shingleID > maxHashCode:
            maxHashCode = shingleID
    return [minHashCode, maxHashCode]

minhash_udf = udf(generate_minhash_signatures, ArrayType(IntegerType()))
df_with_minhash = df.withColumn('min_max_hash', minhash_udf("shingles", coeffA, coeffB))
df_with_minhash.show()

实际udf:

TypeError: Invalid argument, not a string or column: [2856022824, 2966132496, 947839218, 1658426276, 1862779421, 3729685802, 1710806966, 2696513050, 3630333076, 2555745391] of type <class 'list'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

1 个答案:

答案 0 :(得分:1)

您的udf期望所有三个参数均为列。 coeffAcoeffB可能不仅仅是您需要使用lit转换为列对象的数值:

import pyspark.sql.functions as f
df.withColumn('min_max_hash', minhash_udf(f.col("shingles"), f.lit(coeffA), f.lit(coeffB)))

如果coeffAcoeffB是列表,请使用f.array创建以下文字:

df.withColumn('min_max_hash', 
  minhash_udf(f.col("shingles"), 
  f.array(*map(f.lit, coeffA)),
  f.array(*map(f.lit, coeffB))
)

或按以下方式分别分隔列参数和非列参数:

def generate_minhash_signatures(coeffA, coeffB, numHashes)
    def generate_minhash_signatures_inner(shingles):
        signature = []
        for i in range(0, numHashes):
            minHashCode = nextPrime + 1
            for shingleID in shingles:
                hashCode = (coeffA[i] * shingleID + coeffB[i]) % nextPrime

                if hashCode < minHashCode:
                    minHashCode = hashCode

            signature.append(minHashCode)
        return signature
    return f.udf(generate_minhash_signatures_inner, ArrayType(IntegerType()))

然后您可以将函数调用为:

df.withColumn('min_max_hash', generate_minhash_signatures(coeffA, coeffB, numHashes)("shingles"))