Polars, Python, как изменить количество условий, вводимых при создании нового столбца ⇐ Python

[Anonymous]

У меня есть большие наборы данных (от 100 000 до 4 миллионов строк), в которых я ищу различные релевантные коды в нескольких столбцах. Например, если бы я хотел идентифицировать каждую строку, которая начинается со строки «302», я бы сделал:

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

import polars as pl

df = pl.DataFrame({
'Codes_1': ['302E513', '301E513', '302E512'],
'Codes_2': ['303E513', '306E510', '302E512']}).lazy()

conditions = ['302E513', '306E510']
column_names = ['Codes_1', 'Codes_2']

#create new column
df = df.with_columns(
pl.when(pl.any_horizontal(
pl.col(column_names).str.starts_with(conditions[0]),
pl.col(column_names).str.starts_with(conditions[1])))
.then(1.0)
.otherwise(0.0)
.alias('Column_name')
)

Это очень раздражает, когда я ищу, скажем, 4 кода вместо 2, чтобы вводить каждый из кодов, чтобы сформировать новый столбец:

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

import polars as pl

df = pl.DataFrame({
'Codes_1': ['302E513', '301E513', '302E512'],
'Codes_2': ['303E513', '306E510', '302E512']}).lazy()

conditions = ['302E513', '306E510', '5164E23', '302E514']
column_names = ['Codes_1', 'Codes_2']

#create new column
df = df.with_columns(
pl.when(pl.any_horizontal(
#Tedious part
pl.col(column_names).str.starts_with(conditions[0]),
pl.col(column_names).str.starts_with(conditions[1]),
pl.col(column_names).str.starts_with(conditions[2]),
pl.col(column_names).str.starts_with(conditions[3])
))
.then(1.0)
.otherwise(0.0)
.alias('Column_name')
)

Я знаю, что это можно сделать с помощью pandas, обновив маску с помощью цикла for

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

import pandas as pd

df = pd.DataFrame({
'Codes_1': ['302E513', '301E513', '302E512'],
'Codes_2': ['303E513', '306E510', '302E512']})

conditions = ['302E513', '306E510']
column_names = ['Codes_1', 'Codes_2']

#loop to create new column
mask = False
for code in conditions:
mask |= df[column_names].eq(code).any(axis=1)

df['Column_name'] = 0.0
df.loc[mask, 'Column_name'] = 1.0
print(df['Column_name'])

И я мог бы изменить количество условий на любое число, и этот код выполнится. Однако я бы предпочел использовать поляры, поскольку они быстрее и не перегружают оперативную память моей машины для больших наборов данных. Любая помощь приветствуется.

Подробнее здесь: https://stackoverflow.com/questions/792 ... aking-a-ne