Родной API Polars вместо медленного «map_elements» ⇐ Python

[Anonymous]

У меня есть кадр данных с полями «возраст», диагноз и т. д., который я читаю из CSV-файла. Мне нужно добавить код столбца (код заболевания), который зависит от диагноза и/или возраста. В настоящее время я использую метод map_elements:

Код: Выделить всё

disease_codes = {
"malaria": {"under_5": "a", "over_5": "b"},
"PUD": "c",
"Asthma": "d",
}
# The callable to return the appropriate code
def return_code(row):
diagnosis = row["diagnosis"]
age = row["age"]
dx_return = disease_codes.get(diagnosis, "Undefined")
if type(dx_return) == dict:
if age >= 5:
return dx_return.get("over_5")
return dx_return.get("under_5")
return dx_return
df.with_columns(pl.struct(["diagnosis", "age"]).map_elements(return_code).alias("code"))

Все работает как положено. Но поскольку этот метод использует цикл for, в документации есть предупреждение:

Этот метод намного медленнее, чем API собственных выражений. Используйте его только в том случае, если иначе вы не можете реализовать свою логику.

Я пытался найти способ добиться этого с помощью собственного API и подумал, что map_batches подошёл близко:

Код: Выделить всё

def return_code(diagnosis, age):
dx_return = disease_codes.get(diagnosis, "Undefined")
if type(dx_return) == dict:
if age >= 5:
return dx_return.get("over_5")
return dx_return.get("under_5")
return dx_return

df.with_columns(
(pl.struct(["diagnosis", "age"]).map_batches(
lambda x: return_code(x.struct.field("diagnosis"), x.struct.field("age"))
)).alias("code")
)

Очевидно, я не могу этого сделать, так как получаю ошибку:

Код: Выделить всё

TypeError: cannot use `__getitem__` on Series of dtype Struct([Field('diagnosis', Utf8), Field('age', Int64)]) with argument 'diagnosis' of type 'str'

Глядя на документацию, я считаю, что это единственный метод, который поможет. Я застрял на более медленном подходе? Есть ли что-то, что я делаю не так? Что бы вы предложили?
Ожидаемый результат, который я получаю при рабочем подходе, должен быть следующим:



Подробнее здесь: https://stackoverflow.com/questions/774 ... p-elements