Блог

Я создал тестовый набор данных

test_dataset = spark.createDataFrame([tuple(float(i) for i in np.random.rand(61)) for _ in range(10)], [f'col_{i}' for i in range(61)])

У меня есть функция pandas_udf, которая получает оценку от модели — например, LGBMRanker.

 def some_func(row):
  
  return pd.Series(lgb_ranker.predict(np.random.rand(61).reshape(1, 61)))

some_func_udf = pandas_udf(some_func, FloatType())
 

и набор столбцов, которые я хочу передать в модель cols_set = [f'col_{i}' for i in range(61)]

когда я пытаюсь вычислить оценку из этого списка столбцов и добавить ее в фрейм данных в качестве нового столбца test_dataset.select(some_func_udf(test_dataset.select(cols_set))).display()

Я попадаю по ошибке:

 TypeError: Invalid argument, not a string or column: DataFrame[col_0: double, col_1: double,
col_2: double, col_3: double, col_4: double, col_5: double, col_6: double, col_7: double, col_8:
double, col_9: double, col_10: double, col_11: double, col_12: double, col_13: double, col_14:
double, col_15: double, col_16: double, col_17: double, col_18: double, col_19: double, col_20:
double, col_21: double, col_22: double, col_23: double, col_24: double, col_25: double, col_26:
double, col_27: double, col_28: double, col_29: double, col_30: double, col_31: double, col_32:
double, col_33: double, col_34: double, col_35: double, col_36: double, col_37: double, col_38:
double, col_39: double, col_40: double, col_41: double, col_42: double, col_43: double, col_44:
double, col_45: double, col_46: double, col_47: double, col_48: double, col_49: double, col_50:
double, col_51: double, col_52: double, col_53: double, col_54: double, col_55: double, col_56:
double, col_57: double, col_58: double, col_59: double, col_60: double]
of type <class 'pyspark.sql.dataframe.DataFrame'>.
For column literals, use 'lit', 'array', 'struct' or 'create_map' function.
 

Та же ошибка, которую я получаю при попытке выполнить withColumn() выражение:

test_dataset.withColumn('score', some_func_udf(test_dataset.select(cols_set))).display()

Как я могу добавить новый столбец, который является результатом некоторого преобразования (функции или экземпляра модели) в список столбцов?