Вариант 1. Используйте UDF для одного столбца за раз

Самым простым подходом было бы переписать вашу функцию так, чтобы она принимала строку в качестве аргумента (чтобы это было string -> string) и использовала UDF. Есть хороший пример здесь. Это работает на одном столбце за раз. Итак, если ваш DataFrame имеет разумное количество столбцов, вы можете применить UDF к каждому столбцу по одному:

from pyspark.sql.functions import col
new_df = df.select(udf(col("col1")), udf(col("col2")), ...)

Пример

df = sc.parallelize([[1, 4], [2,5], [3,6]]).toDF(["col1", "col2"])
df.show()
+----+----+
|col1|col2|
+----+----+
|   1|   4|
|   2|   5|
|   3|   6|
+----+----+

def plus1_udf(x):
    return x + 1
plus1 = spark.udf.register("plus1", plus1_udf)

new_df = df.select(plus1(col("col1")), plus1(col("col2")))
new_df.show()
+-----------+-----------+
|plus1(col1)|plus1(col2)|
+-----------+-----------+
|          2|          5|
|          3|          6|
|          4|          7|
+-----------+-----------+

Вариант 2: Сопоставьте весь DataFrame сразу

map доступен для Scala DataFrame, но на данный момент недоступен в PySpark. API нижнего уровня СДР имеет функцию map в PySpark. Итак, если у вас слишком много столбцов для преобразования по одному, вы можете работать с каждой отдельной ячейкой в ​​DataFrame следующим образом:

def map_fn(row):
    return [api_function(x) for (column, x) in row.asDict().items()

column_names = df.columns
new_df = df.rdd.map(map_fn).toDF(df.columns)

Пример

df = sc.parallelize([[1, 4], [2,5], [3,6]]).toDF(["col1", "col2"])
def map_fn(row):
   return [value + 1 for (_, value) in row.asDict().items()]

columns = df.columns
new_df = df.rdd.map(map_fn).toDF(columns)
new_df.show()
+----+----+
|col1|col2|
+----+----+
|   2|   5|
|   3|   6|
|   4|   7|
+----+----+

Контекст

документацияforeach дает только пример печати, но мы можем убедиться, глядя на код, что он действительно ничего не возвращает.

Вы можете прочитать о pandas_udf в эта почта, но кажется, что он больше всего подходит для векторизованных функций, которые, как вы указали, вы не можете использовать из-за api_function.

Решение:

udf_func = udf(func, StringType())
for col_name in spark_df.columns:
    spark_df = spark_df.withColumn(col_name, udf_func(lit(col_name), col_name))
return spark_df.toPandas()

Есть 3 ключевых момента, которые помогли мне понять это:

1. Если вы используете withColumn с именем существующего столбца (col_name), Spark «перезаписывает»/затеняет исходный столбец. По сути, это создает видимость редактирования столбца напрямую, как если бы он был изменяемым.
2. Создавая цикл для исходных столбцов и повторно используя одну и ту же переменную DataFrame spark_df, я использую тот же принцип для имитации изменяемого DataFrame, создавая цепочку преобразований по столбцам, каждый раз «перезаписывая» столбец (согласно # 1 - см. ниже )
3. Spark UDFs ожидает, что все параметры будут типами Column, что означает, что он пытается разрешить значения столбца для каждого параметра. Поскольку первый параметр api_function — это буквальное значение, которое будет одинаковым для всех строк в векторе, вы должны использовать функцию lit(). Простая передача col_name в функцию попытается извлечь значения столбца для этого столбца. Насколько я мог судить, передача col_name эквивалентна передаче col(col_name).

Предполагая 3 столбца «a», «b» и «c», развертывание этой концепции будет выглядеть так:

spark_df = spark_df.withColumn('a', udf_func(lit('a'), 'a')
                   .withColumn('b', udf_func(lit('b'), 'b')
                   .withColumn('c', udf_func(lit('c'), 'c')