如何通过PySpark的isin方法基于另一DataFrame筛选数据

问题描述 #

我有数据量较大的df1，想要筛选出其中id存在于df2中的数据，编写了如下代码：

df1.filter(col('id').isin(df2.select('id')))

运行后出现报错：

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/pyspark/sql/column.py in isin(self, *cols)
    441         if len(cols) == 1 and isinstance(cols[0], (list, set)):
    442             cols = cols[0]
--> 443         cols = [c._jc if isinstance(c, Column) else _create_column_from_literal(c) for c in cols]
    444         sc = SparkContext._active_spark_context
    445         jc = getattr(self._jc, "isin")(_to_seq(sc, cols))

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/pyspark/sql/column.py in <listcomp>(.0)
    441         if len(cols) == 1 and isinstance(cols[0], (list, set)):
    442             cols = cols[0]
--> 443         cols = [c._jc if isinstance(c, Column) else _create_column_from_literal(c) for c in cols]
    444         sc = SparkContext._active_spark_context
    445         jc = getattr(self._jc, "isin")(_to_seq(sc, cols))

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/pyspark/sql/column.py in _create_column_from_literal(literal)
     33 def _create_column_from_literal(literal):
     34     sc = SparkContext._active_spark_context
---> 35     return sc._jvm.functions.lit(literal)
     36 
     37 

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/lib/py4j-0.10.7-src.zip/py4j/java_gateway.py in __call__(self, *args)
   1246 
   1247     def __call__(self, *args):
-> 1248         args_command, temp_args = self._build_args(*args)
   1249 
   1250         command = proto.CALL_COMMAND_NAME +\

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/lib/py4j-0.10.7-src.zip/py4j/java_gateway.py in _build_args(self, *args)
   1216 
   1217         args_command = "".join(
-> 1218             [get_command_part(arg, self.pool) for arg in new_args])
   1219 
   1220         return args_command, temp_args

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/lib/py4j-0.10.7-src.zip/py4j/java_gateway.py in <listcomp>(.0)
   1216 
   1217         args_command = "".join(
-> 1218             [get_command_part(arg, self.pool) for arg in new_args])
   1219 
   1220         return args_command, temp_args

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/lib/py4j-0.10.7-src.zip/py4j/protocol.py in get_command_part(parameter, python_proxy_pool)
    296             command_part += ";" + interface
    297     else:
--> 298         command_part = REFERENCE_TYPE + parameter._get_object_id()
    299 
    300     command_part += "\n"

/opt/cloudera/parcels/CDH-7.1.3-1.cdh7.1.3.p0.4992530/lib/spark/python/pyspark/sql/dataframe.py in __getattr__(self, name)
   1302         if name not in self.columns:
   1303             raise AttributeError(
-> 1304                 "'%s' object has no attribute '%s'" % (self.__class__.__name__, name))
   1305         jc = self._jdf.apply(name)
   1306         return Column(jc)

AttributeError: 'DataFrame' object has no attribute '_get_object_id'

错误原因 #

isin()方法接收的是列表、集合或单个值，但你传入的df2.select('id')返回的是一个DataFrame对象，而非可迭代的id值集合。Spark尝试将这个DataFrame当作字面量处理时，触发了对象属性不存在的错误。

解决方法 #

方法1：将df2的id转为列表（适合df2数据量较小的场景） #

如果df2的id数量不多，可以先把id提取成Python列表，再传入isin()：

from pyspark.sql.functions import col

# 提取df2的id为列表
id_list = df2.select('id').rdd.flatMap(lambda x: x).collect()
# 筛选df1
filtered_df = df1.filter(col('id').isin(id_list))

注意：collect()会把数据拉到Driver节点内存，若df2数据量过大，可能导致内存溢出，这种情况不推荐用此方法。

方法2：使用JOIN或EXISTS子查询（适合大数据量场景） #

针对大数据量的df1和df2，更高效的方式是用Spark的分布式操作，比如内连接或者exists子查询：

方式A：内连接（Inner Join）

# 内连接后保留df1的所有列
filtered_df = df1.join(df2.select('id'), on='id', how='inner')

如果担心df2有重复id导致结果重复，可以先对df2去重：

filtered_df = df1.join(df2.select('id').distinct(), on='id', how='inner')

方式B：EXISTS子查询

用Spark SQL的exists语法实现：

# 注册临时表
df1.createOrReplaceTempView("df1")
df2.createOrReplaceTempView("df2")

# 执行SQL查询
filtered_df = spark.sql("""
    SELECT * FROM df1
    WHERE EXISTS (
        SELECT 1 FROM df2 WHERE df2.id = df1.id
    )
""")

这种方式和join性能相近，都是分布式执行，不会把数据拉到Driver节点，适合大数据量场景。

内容的提问来源于stack exchange，提问作者Nabih Bawazir