Skip to main content
 首页 » 编程设计

python之如何检查 Pyspark Map 中是否存在键或值

2024年10月01日39cloudgamer

我在 spark DF 中有一个 Map 列,我想根据特定键过滤此列(即,如果 map 中的键与所需值匹配,则保留该行)。

例如,我的架构定义为:

df_schema = StructType( 
  [StructField('id', StringType()), 
   StructField('rank', MapType(StringType(), IntegerType()))] 
) 

我的示例数据是:

{ "id": "0981850006", "rank": {"a": 1} } 

有什么方法可以在不使用 explode() 的情况下过滤“a”在“等级”中的行上的 df 吗?

对于给定的 json,是否有比我定义的更好的模式表示?

请您参考如下方法:

使用 rank.key 访问 key 意味着 rank 是一个 StructType()。尽管爆炸可能是最好的解决方案,但让我们构建一个 UDF 来评估 k 是否是 rank 的键。

首先让我们创建我们的数据框:

from pyspark.sql.types import * 
df_schema = StructType( 
  [StructField('id', StringType()), 
   StructField('rank', MapType(StringType(), IntegerType()))] 
) 
df = spark.createDataFrame([ 
    ["0981850006", {"a": 1}],  
    ["0981850006", {"b": 2, "c": 3}],  
], df_schema) 

现在我们的 UDF:

def isKey(k,d): 
    return k in d.keys() 
 
isKey_udf = lambda k: psf.udf(lambda d: isKey(k,d), BooleanType()) 

给出:

df.withColumn( 
    "is_key",  
    isKey_udf('a')(df.rank) 
) 
    +----------+-------------------+------+ 
    |        id|               rank|is_key| 
    +----------+-------------------+------+ 
    |0981850006|        Map(a -> 1)|  true| 
    |0981850006|Map(b -> 2, c -> 3)| false| 
    +----------+-------------------+------+