Как применить результат функции Python к новому столбцу в фрейме данных Polars
Я пробовал искать, но многие ответы кажутся устаревшими и больше не действительны для текущей версии Polars. Как применить возвращаемый результат функции Python к каждой строке кадра данных Polars? Я хочу передать в функцию всю строку вместо передачи определенных столбцов.
import polars as pl
def test2(auth, row):
c = row["Group"]
d = row["Val"]
return "{}-{}-{}".format(c, str(d), auth)
df = pl.DataFrame({
'Group': ['A', 'B', 'C', 'D', 'E'],
'Val': [1001, 1002, 1003, 1004, 1005]
})
auth_token = "xxxxxxxxx"
df = df.with_columns(
pl.struct(pl.all())
.map_batches(lambda x: test2(auth_token, x))
.alias("response")
)
print(df)
Приведенный выше код вызывает эту ошибку. Я не понимаю это сообщение. Где мне следует установить strict=False и зачем это нужно?
Traceback (most recent call last):
File "c:\Scripting\Python\Development\Test.py", line 29, in <module>
df = df.with_columns(
File "c:\Scripting\Python\Development\venv\lib\site-packages\polars\dataframe\frame.py", line 8763, in with_columns
return self.lazy().with_columns(*exprs, **named_exprs).collect(_eager=True)
File "c:\Scripting\Python\Development\venv\lib\site-packages\polars\lazyframe\frame.py", line 1942, in collect
return wrap_df(ldf.collect(callback))
polars.exceptions.ComputeError: TypeError: unexpected value while building Series of type Int64; found value of type String: "C"
Hint: Try setting `strict=False` to allow passing data with mixed types.
Я знаю, что мог бы сделать это, указав конкретные столбцы, такие как приведенный ниже код, но я хочу передать всю строку, а затем выбрать, какие столбцы будут использоваться внутри функции. Любая помощь будет оценена по достоинству. Спасибо.
df = df.with_columns(
(
pl.struct(["Group", "Val"]).map_batches(
lambda x: test(auth_token, x.struct.field("Group"), x.struct.field("Val"))
)
).alias("api_response")
)
Ответы 1
Вместо этого вы можете использовать map_elements().
df = df.with_columns(
pl.struct(pl.all())
.map_elements(lambda x: test2(auth_token, x))
.alias("response")
)
┌───────┬──────┬──────────────────┐
│ Group ┆ Val ┆ response │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str │
╞═══════╪══════╪══════════════════╡
│ A ┆ 1001 ┆ A-1001-xxxxxxxxx │
│ B ┆ 1002 ┆ B-1002-xxxxxxxxx │
│ C ┆ 1003 ┆ C-1003-xxxxxxxxx │
│ D ┆ 1004 ┆ D-1004-xxxxxxxxx │
│ E ┆ 1005 ┆ E-1005-xxxxxxxxx │
└───────┴──────┴──────────────────┘
Если вы хотите разделить результат функции на несколько столбцов (при условии, что вы можете изменить функцию):
- используя словарь в качестве возвращаемого значения, поэтому в конечном итоге он станет столбцом struct в DataFrame.
- unnest() для разделения полученной структуры на отдельные столбцы
def test2(auth, row):
c = row["Group"]
d = row["Val"]
return {"c": c, "d": str(d), "auth": auth}
(
df
.with_columns(
pl.struct(pl.all())
.map_elements(lambda x: test2(auth_token, x))
.alias("response")
)
).unnest("response")
┌───────┬──────┬─────┬──────┬───────────┐
│ Group ┆ Val ┆ c ┆ d ┆ auth │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str ┆ str ┆ str │
╞═══════╪══════╪═════╪══════╪═══════════╡
│ A ┆ 1001 ┆ A ┆ 1001 ┆ xxxxxxxxx │
│ B ┆ 1002 ┆ B ┆ 1002 ┆ xxxxxxxxx │
│ C ┆ 1003 ┆ C ┆ 1003 ┆ xxxxxxxxx │
│ D ┆ 1004 ┆ D ┆ 1004 ┆ xxxxxxxxx │
│ E ┆ 1005 ┆ E ┆ 1005 ┆ xxxxxxxxx │
└───────┴──────┴─────┴──────┴───────────┘
Это сработало отлично. Спасибо. Я неправильно понял разницу между Map_batches и Map_elements.
Кстати, знаете ли вы, как разбить ответ функции на отдельные столбцы? Например, вместо того, чтобы объединять значения в одну строку и возвращать их в качестве ответа, я хотел бы вернуть каждое значение (c, d и auth) в отдельные столбцы в фрейме данных? Спасибо.
вы можете вернуть кортеж или список из функции, а затем разбить список на столбцы
Есть ли у вас краткий пример синтаксиса для этого? Если да, то очень признателен. Спасибо!
Обновил ответ
в вашем примере вы можете просто использовать
map_elements()вместоmap_batches(). Но это будет довольно медленно, потому что вы не используете никакую оптимизацию Polars таким образом.