如何使用PySpark DataFrame选择给定元组列表的行？

Question

假设我们有DataFrame这样：

+--------+--------------+-----+--------------------+
|aid     |bid           |value|                time|
+--------+--------------+-----+--------------------+
|       1|             1| 81.0|2006-08-25 14:13:...|
|       1|             1| 81.0|2006-08-25 14:27:...|
|       1|             2| 81.0|2006-08-25 14:56:...|
|       1|             2| 81.0|2006-08-25 15:00:...|
|       1|             3| 81.0|2006-08-25 15:31:...|
|       1|             3| 81.0|2006-08-25 15:38:...|
|       1|             4|  0.0|2006-08-30 11:59:...|
|       1|             4|  0.0|2006-08-30 13:59:...|
|       2|             1|  0.0|2006-08-30 12:11:...|
|       2|             1|  0.0|2006-08-30 14:13:...|
|       2|             2|  0.0|2006-08-30 12:30:...|
|       2|             2|  0.0|2006-08-30 14:30:...|
|       2|             3|  0.0|2006-09-05 12:29:...|
|       2|             3|  0.0|2006-09-05 14:31:...|
|       3|             1|  0.0|2006-09-05 12:42:...|
|       3|             1|  0.0|2006-09-05 14:43:...|
+--------+--------------+-----+--------------------+

我知道我可以这样做：

df_data.where(col('bid')
       .isin([1,2,3])).show()

以便仅选择bid之一[1,2,3]的行。

但是，我希望能够根据[(1,1), (2,2), (3,1)]和aid两列的元组bid列表选择子集。

所以基本上“类似”

df_data.where(col(['aid', 'bid'])
       .isin([(1,1), (2,2), (3,1)])).show()

有办法做到这一点吗？

我可以想象这样的事情：

sql.sql('SELECT * FROM df_data WHERE (scope_id, measurement_id) IN ((1,1))')

但这将抛出：

AnalysisException: "cannot resolve '(struct(df_data.`aid`, df_data.`bid`) IN (struct(1, 1)))' due to data type mismatch: Arguments must be same type; line 1 pos 55"

Answer 1

 

我可以想到三种方式。

方法1：使用reduce帮助检查所有条件

伪代码(s, m) IN [(1,1), (2,2), (3,1)]相当于：

(s == 1 and m == 1) or (s == 2 and m == 2) or (s == 3 and m == 3)

您可以使用列表推导和reduce检查所有这些条件。

import pyspark.sql.functions as f
check_list = [(1,1), (2,2), (3,1)]
df.where(
        reduce(
            lambda u, v: u|v,
            [(f.col("aid") == x) & (f.col("bid") == y) for (x,y) in check_list]
        )
    )\
    .select("aid", "bid", "value")\
    .show()
#+---+---+-----+
#|aid|bid|value|
#+---+---+-----+
#|  1|  1| 81.0|
#|  1|  1| 81.0|
#|  2|  2|  0.0|
#|  2|  2|  0.0|
#|  3|  1|  0.0|
#|  3|  1|  0.0|
#+---+---+-----+

方法2：将ID连接为字符串

创建一个临时列作为两个id列的字符串连接。然后检查该字符串是否与字符串列表匹配。

check_list = [(1,1), (2,2), (3,1)]
check_list_str = [",".join([str(x) for x in item]) for item in check_list]

df.withColumn("combined_id", f.concat(f.col("aid"), f.lit(","), f.col("bid")))\
    .where(f.col("combined_id").isin(check_list_str))\
    .select("aid", "bid", "value")\
    .show()
#+---+---+-----+
#|aid|bid|value|
#+---+---+-----+
#|  1|  1| 81.0|
#|  1|  1| 81.0|
#|  2|  2|  0.0|
#|  2|  2|  0.0|
#|  3|  1|  0.0|
#|  3|  1|  0.0|
#+---+---+-----+

方法3：使用UDF

创建udf以检查布尔条件。

check_list = [(1,1), (2,2), (3,1)]
check_id_isin = f.udf(lambda x, y: (x, y) in check_list, BooleanType())

df.where(check_id_isin(f.col("aid"), f.col("bid")) == True)\
    .select("aid", "bid", "value")\
    .show()
#+---+---+-----+
#|aid|bid|value|
#+---+---+-----+
#|  1|  1| 81.0|
#|  1|  1| 81.0|
#|  2|  2|  0.0|
#|  2|  2|  0.0|
#|  3|  1|  0.0|
#|  3|  1|  0.0|
#+---+---+-----+

编辑作为@StefanFalk pointed out，可以将udf更一般地写为：

check_id_isin = f.udf(lambda *idx: idx in check_list, BooleanType())

这将允许可变数量的输入参数。

Answer 2

另一种选择是

[DllImport("gpc.dll", CharSet = CharSet.Ansi, CallingConvention = CallingConvention.StdCall, SetLastError = true)]
    private static extern void _gpc_create_polygon([In, Out] ref _gpc_vertex origin, [In, Out] ref int sides, [In, Out] ref double radius, 
                                                  [In, Out] ref _gpc_polygon result);