时间戳

首页 > 解决方案 > 从 df 列的列表中过滤期望值

python - 从 df 列的列表中过滤期望值

问题描述

我有一个包含以下列的数据框：

raw_col
['a','b','c']
['b']
['a','b']
['c']

我想根据条件语句返回具有单个值的列。我写了以下函数：

def filter_func(elements):
  if "a" in elements:
    return "a"
  else:
    return "Other"

在列上运行函数时df.withColumn("col", filter_func("raw_col"))出现以下错误col should be Column

这里有什么问题？我应该怎么办？

标签： pythonpandasapache-sparkpyspark

解决方案

您可以使用array_contains功能：

import pyspark.sql.functions as f

df = df.withColumn("col", f.when(f.array_contains("raw_col", f.lit("a")), f.lit("a")).otherwise(f.lit("Other")))

但是如果你有一个复杂的逻辑并且需要使用filter_func，则需要创建一个 UDF：

@f.udf()
def filter_func(elements):
    if "a" in elements:
        return "a"
    else:
        return "Other"

df = df.withColumn("col", filter_func("raw_col"))

推荐阅读