具有可选"and"条件的 Python 函数

Question

我想弄清楚如果它有意义将下面2个相似的函数合并为1个(这是从实用的角度考虑的——哪个更容易处理)。如果是，那么how

?

• get_result_1():始终需要
• get_result_2():get_result_1()+ "and"包含变量column_name &input_value

有问题的函数:

def get_result_1(item: str, type: str) -> Column:
c = ((f.col("item") == item) & (f.col("type") == type))
return c
def get_result_2(item: str, type: str, input_1: str, input_2: str) -> Column:
c = (get_result_1(item, type)
& (f.col(input_1) == input_2) #<-- make this part of get_result_1?
)
return c

用例:

output_df = (
df
.withColumn("category", 
f.when(get_result_1("A", "xyz1"), f.lit("C1"))
.when(get_result_2("B", "xyz2", "level1", "high"), f.lit("C2"))
.when(get_result_2("B", "xyz3", "level2", "L7"), f.lit("C3"))
)
)

我一直在探索各种参数，但我不能使它工作。在这一点上，我甚至怀疑我是否走在正确的道路上。

Answer 1

从Spark端来看，这应该无关紧要，因为无论哪种方式都应该产生相同的查询。您可以使用以下两种方法对df.explain()进行检查。

您可以通过将input_1和input_2设置为可选并检查它们是否通过来组合这两个函数。如果您愿意，还可以加一些额外的检查。

def get_result(item: str, item_type: str, input_1: str = None, input_2: str = None) -> Column:
c = ((f.col("item") == item) & (f.col("type") == item_type))
if input_1:
c = c & (f.col(input_1) == input_2)
return c