[解決済み] TypeError: 'Column' オブジェクトは WithColumn を使用して呼び出すことができません。

質問

データフレーム(df)に新しいカラムを追加したいのですが、関数 get_distance :

def get_distance(x, y):
    dfDistPerc = hiveContext.sql("select column3 as column3, \
                                  from tab \
                                  where column1 = '" + x + "' \
                                  and column2 = " + y + " \
                                  limit 1")

    result = dfDistPerc.select("column3").take(1)
    return result

df = df.withColumn(
    "distance",
    lit(get_distance(df["column1"], df["column2"]))
)

でも、こうなるんです。

TypeError: 'Column' object is not callable

xとyがあるから起こるのだと思います。 Column オブジェクトに変換される必要があり、私は String をクエリで使用することができます。正しいですか？もしそうなら、どうすればいいのでしょうか？

どのように解決するのですか？

Sparkは、使用している関数が通常の関数ではなく、UDFであることを知っている必要があります。

つまり、データフレームでUDFを使うには、2つの方法があるのです。

方法-1: @udfアノテーションを使用する

@udf
def get_distance(x, y):
    dfDistPerc = hiveContext.sql("select column3 as column3, \
                                  from tab \
                                  where column1 = '" + x + "' \
                                  and column2 = " + y + " \
                                  limit 1")

    result = dfDistPerc.select("column3").take(1)
    return result

df = df.withColumn(
    "distance",
    lit(get_distance(df["column1"], df["column2"]))
)

方法-2: pyspark.sql.functions.udf を使用して udf を再定義します。

def get_distance(x, y):
    dfDistPerc = hiveContext.sql("select column3 as column3, \
                                  from tab \
                                  where column1 = '" + x + "' \
                                  and column2 = " + y + " \
                                  limit 1")

    result = dfDistPerc.select("column3").take(1)
    return result

calculate_distance_udf = udf(get_distance, IntegerType())

df = df.withColumn(
    "distance",
    lit(calculate_distance_udf(df["column1"], df["column2"]))
)