Эффективный метод добавления производных данных в Polars multiIndex LazyFrame

Я работаю с Polars, и мне нужно добавить производные данные в многоиндексный LazyFrame. Чтобы изменить данные, я поворачиваю DataFrame, не выполняя каких-либо агрегаций, выполняя некоторые вычисления, а затем возвращая DataFrame обратно в исходный формат. Эту операцию необходимо выполнить на разных уровнях индекса. В документации Polars упоминается, что операции поворота недоступны в ленивом режиме. Из-за поворота/отмены поворота необходимо переключаться между нетерпеливым и ленивым режимами. Есть ли более эффективный способ добиться этого без переключения между нетерпеливым и ленивым режимами?

Вот пример:

import numpy as np
import polars as pl

def get_raw_data() -> pl.LazyFrame:
    """Generate random a multiindex LazyFrame with different size of indexes as example."""
    names = np.array(['A', 'B', 'C'])
    measures = np.array(['height', 'width'])
    repeats: np.ndarray = np.array([3, 3, 2])
    return pl.DataFrame({
        'id': np.repeat(names, repeats*measures.size),
        'measure': np.concatenate([np.repeat(measures, rep) for rep in repeats]),
        'date': np.concatenate([np.arange(size) for size in np.repeat(repeats, measures.size)]),
        'value': np.round(np.random.default_rng(111).random(measures.size*repeats.sum()), 2),
    }).lazy()

print(get_raw_data().collect())
# shape: (16, 4)
# ┌─────┬─────────┬──────┬───────┐
# │ id  ┆ measure ┆ date ┆ value │
# │ --- ┆ ---     ┆ ---  ┆ ---   │
# │ str ┆ str     ┆ i64  ┆ f64   │
# ╞═════╪═════════╪══════╪═══════╡
# │ A   ┆ height  ┆ 0    ┆ 0.15  │
# │ A   ┆ height  ┆ 1    ┆ 0.17  │
# │ A   ┆ height  ┆ 2    ┆ 0.51  │
# │ A   ┆ width   ┆ 0    ┆ 0.66  │
# │ A   ┆ width   ┆ 1    ┆ 0.77  │
# │ …   ┆ …       ┆ …    ┆ …     │
# │ B   ┆ width   ┆ 2    ┆ 0.72  │
# │ C   ┆ height  ┆ 0    ┆ 0.08  │
# │ C   ┆ height  ┆ 1    ┆ 0.42  │
# │ C   ┆ width   ┆ 0    ┆ 0.4   │
# │ C   ┆ width   ┆ 1    ┆ 0.94  │
# └─────┴─────────┴──────┴───────┘

def expr_add_categories() -> pl.Expr:
    """Generate a sample list of expressions to add some derived categories."""
    return [(pl.col('height')/pl.col('width')).alias('ratio')]

def expr_add_ids() -> pl.Expr:
    """Generate a sample list of expressions to add some derived ids."""
    return [
        (pl.col('A') / pl.col('B')).alias('AB'),
        (pl.col('A') / pl.col('C')).alias('AC')
    ]

def add_categories(df: pl.LazyFrame) -> pl.LazyFrame:
    """Add various derived categories to LazyFrame."""
    return (
        df
        .collect()  # pivot requires eager mode
        .pivot(index=['id', 'date'], columns='measure', values='value')
        .lazy()     # back to lazy mode
        .with_columns(expr_add_categories())
        .melt(id_vars=['id', 'date'], variable_name='measure')
        .drop_nulls()
        .select(['id', 'measure', 'date', 'value'])
        .sort(['id', 'measure', 'date'])
        .set_sorted(['id', 'measure', 'date'])
    )

def add_ids(df: pl.LazyFrame) -> pl.LazyFrame:
    """Add various derived IDs to LazyFrame."""
    return (
        df
        .collect()  # pivot requires eager mode
        .pivot(index=['measure', 'date'], columns='id', values='value')
        .lazy()     # back to lazy mode
        .with_columns(expr_add_ids())
        .melt(id_vars=['measure', 'date'], variable_name='id')
        .drop_nulls()
        .select(['id', 'measure', 'date', 'value'])
        .sort(['id', 'measure', 'date'])
        .set_sorted(['id', 'measure', 'date'])
    )

def get_modified_data() -> pl.LazyFrame:
    """Get raw data and add derived categories and names to LazyFrame."""
    return (
        get_raw_data()
        .pipe(add_categories)
        .pipe(add_ids)
    )

print(get_modified_data().collect())
# shape: (39, 4)
# ┌─────┬─────────┬──────┬──────────┐
# │ id  ┆ measure ┆ date ┆ value    │
# │ --- ┆ ---     ┆ ---  ┆ ---      │
# │ str ┆ str     ┆ i64  ┆ f64      │
# ╞═════╪═════════╪══════╪══════════╡
# │ A   ┆ height  ┆ 0    ┆ 0.15     │
# │ A   ┆ height  ┆ 1    ┆ 0.17     │
# │ A   ┆ height  ┆ 2    ┆ 0.51     │
# │ A   ┆ ratio   ┆ 0    ┆ 0.227273 │
# │ A   ┆ ratio   ┆ 1    ┆ 0.220779 │
# │ …   ┆ …       ┆ …    ┆ …        │
# │ C   ┆ height  ┆ 1    ┆ 0.42     │
# │ C   ┆ ratio   ┆ 0    ┆ 0.2      │
# │ C   ┆ ratio   ┆ 1    ┆ 0.446809 │
# │ C   ┆ width   ┆ 0    ┆ 0.4      │
# │ C   ┆ width   ┆ 1    ┆ 0.94     │
# └─────┴─────────┴──────┴──────────┘

# *************************************************************
# Python: 3.12.0
# Numpy: 1.26.4
# Polars: 0.20.31
# *************************************************************

Редактировать: Предположим, что содержимое DataFrame является динамическим. Уникальность элементов любого из уровней индекса заранее неизвестна.

Редактировать: Вот пример «lazypivot» в виде UDF. Я не уверен, как применить функцию агрегирования без использования лямбда-функции. Лучшие предложения приветствуются.

import warnings
from typing import Callable
import polars as pl

def lazypivot(df: pl.LazyFrame,
              index: str | list[str] | None,
              columns: str | list[str] | None,
              values: str | list[str] | None,
              column_values: str | list[str] | None = None,
              aggregate_function: Callable | None = None,
              maintain_order: bool = True,
              sort_columns: bool = True,
              ) -> pl.LazyFrame:
    """Pivot a LazyFrame with or without aggregation."""

    # Collect unique column values if not provided
    if column_values is None:
        warnings.warn(
            'No column_values provided. Switching between eager and lazy mode necessary to collect unique column values.',
            UserWarning
        )
        collected_df = df.collect()
        column_values = collected_df[columns].unique().sort() if sort_columns else collected_df[columns].unique()
        df = collected_df.lazy()

    # Define the aggregation function
    if aggregate_function is None:
        agg_expr = [pl.col(values).filter(pl.col(columns) == value).first().alias(value) for value in column_values]
    else:
        agg_expr = [aggregate_function(pl.col(values).filter(pl.col(columns) == value)).alias(value) for value in column_values]

    # Perform the pivot
    return df.group_by(index, maintain_order=maintain_order).agg(agg_expr)

df = pl.DataFrame(
    {
        "idx": ["A", "A", "A", "B", "B"],
        "cat": ["x", "y", "z", "x", "y"],
        "val": [1, 2, 3, 4, 5],
    }
)
print(df)

print('pivot with column_values:')
df_new = df.lazy().pipe(lazypivot, index = "idx", columns = "cat", values = "val", column_values=['x', 'y', 'z']).collect()
print(df_new)

print('pivot without column_values:')
df_new = df.lazy().pipe(lazypivot, index = "idx", columns = "cat", values = "val").collect()
print(df_new)

# shape: (5, 3)
# ┌─────┬─────┬─────┐
# │ idx ┆ cat ┆ val │
# │ --- ┆ --- ┆ --- │
# │ str ┆ str ┆ i64 │
# ╞═════╪═════╪═════╡
# │ A   ┆ x   ┆ 1   │
# │ A   ┆ y   ┆ 2   │
# │ A   ┆ z   ┆ 3   │
# │ B   ┆ x   ┆ 4   │
# │ B   ┆ y   ┆ 5   │
# └─────┴─────┴─────┘

# pivot with column_values:
# shape: (2, 4)
# ┌─────┬─────┬─────┬──────┐
# │ idx ┆ x   ┆ y   ┆ z    │
# │ --- ┆ --- ┆ --- ┆ ---  │
# │ str ┆ i64 ┆ i64 ┆ i64  │
# ╞═════╪═════╪═════╪══════╡
# │ A   ┆ 1   ┆ 2   ┆ 3    │
# │ B   ┆ 4   ┆ 5   ┆ null │
# └─────┴─────┴─────┴──────┘

# pivot without column_values:
# shape: (2, 4)
# ┌─────┬─────┬─────┬──────┐
# │ idx ┆ x   ┆ y   ┆ z    │
# │ --- ┆ --- ┆ --- ┆ ---  │
# │ str ┆ i64 ┆ i64 ┆ i64  │
# ╞═════╪═════╪═════╪══════╡
# │ A   ┆ 1   ┆ 2   ┆ 3    │
# │ B   ┆ 4   ┆ 5   ┆ null │
# └─────┴─────┴─────┴──────┘
# UserWarning: No column_values provided. Switching between eager and # lazy mode necessary to collect unique column values.