Я хочу удалить одно поле из структуры, сейчас я настроил его так, но есть ли более простой способ добиться этого?
import polars as pl
import polars.selectors as cs
def remove_one_field(df: pl.DataFrame) -> pl.DataFrame:
meta_data_columns = (df.select('meta_data')
.unnest('meta_data')
.select(cs.all() - cs.by_name('system_data')).columns)
print(meta_data_columns)
return (df.unnest('meta_data')
.select(cs.all() - cs.by_name('system_data'))
.with_columns(meta_data=pl.struct(meta_data_columns))
.drop(meta_data_columns))
# Example usage
input_df = pl.DataFrame({
"id": [1, 2],
"meta_data": [{"system_data": "to_remove", "user_data": "keep"}, {"user_data": "keep_"}]
})
output_df = remove_one_field(input_df)
print(output_df)
['user_data']
shape: (2, 2)
┌─────┬───────────┐
│ id ┆ meta_data │
│ --- ┆ --- │
│ i64 ┆ struct[1] │
╞═════╪═══════════╡
│ 1 ┆ {"keep"} │
│ 2 ┆ {"keep_"} │
└─────┴───────────┘
Что-то вроде select
в полях внутри структуры?
Вы можете использовать struct.field() , которая может принимать либо список строк, либо несколько строковых аргументов. Вы знаете свой DataFrame' Schema(), поэтому вы можете легко создать список нужных вам полей.
fields = [c[0] for c in input_df.schema["meta_data"] if c[0] != "system_data"]
input_df.with_columns(
meta_data = pl.struct(
pl.col.meta_data.struct.field(fields)
)
)
┌─────┬───────────┐
│ id ┆ meta_data │
│ --- ┆ --- │
│ i64 ┆ struct[1] │
╞═════╪═══════════╡
│ 1 ┆ {"keep"} │
│ 2 ┆ {"keep_"} │
└─────┴───────────┘
В зависимости от того, насколько динамично вам нужно это делать, это может варьироваться от
# recreate the meta_data column, but only with the user_data field
input_df.with_columns(meta_data=pl.struct(pl.col("meta_data").struct["user_data"]))
который удовлетворяет вашему примеру, что-то вроде
def drop_struct_fields(
df: pl.DataFrame,
struct_column_name: str,
struct_fields_to_drop: str | set,
) -> pl.DataFrame:
if isinstance(struct_fields_to_drop, str):
struct_fields_to_drop = {struct_fields_to_drop}
# Get the struct fields from the column except those we want to drop
output_struct_cols = (
c[0] for c in df.schema[struct_column_name] if c[0] not in struct_fields_to_drop
)
# select only the struct columns we want to include in the output
return df.with_columns(
pl.struct(pl.col(struct_column_name).struct.field(*output_struct_cols)).alias(
struct_column_name
)
)
drop_struct_fields(input_df, "meta_data", "system_data")
# or drop_struct_fields(input_df, "meta_data", {"system_data", "other_struct_field"})
который гораздо более динамичен и, как правило, пригоден для повторного использования.
Отказ от ответственности: output_struct_cols
вдохновлен/украден у Романа Пекара в другом ответе
Честно говоря, я не думаю, что pl.Expr.struct.drop
тоже сойдет с пути как запрос на функцию. Уже существует struct.with_fields
, аналог with_columns
для структур. На Github есть запрос здесь, если вы хотите его поддержать.
Я думаю, у нас также должен быть полярный способ получения полей, возможно, с использованием селекторов.