Общее решение для передачи имен выражений с селекторами в полярах ⇐ Python

[Anonymous]

В полярах (1.7.1) передача имен выражений с помощью различных операций не является простой задачей.
Приведенное ниже связано с проблемой, выявленной в разделе Как получить доступ к имени столбца в полярном выражении? однако я чувствую, что оно настолько отличается, что требует разделения.
В более общем плане проблема заключается в том, что мы хотим применить выражения, включающие несколько столбцов (строго, несколько входных выражений), где мы хотим убедиться, что имя, которое мы получим в конечном итоге, сохраняется на основе «привязочного» выражения, т. е. того, для которого вызывается метод (

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

self._lhs_expr

ниже).
Ниже приведен отдельный пример, иллюстрирующий это.

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

import traceback
import polars as pl

@pl.api.register_expr_namespace("new")
class NewExtensions:
_lhs_expr: pl.Expr

def __init__(self, lhs_expr: pl.Expr):
# Capture the expression on which the methods are to be called.
self._lhs_expr = lhs_expr

def reorder_1(self, weight: pl.Expr) -> pl.Expr:
# This does not seem to work correctly, despite the superficial expectation it would.
return weight.sort_by(self._lhs_expr).alias(self._lhs_expr.meta.output_name())

def reorder_2(self, weight: pl.Expr) -> pl.Expr:
# This "works". It is a kludge that may have unintended consequences, and is not necessarily generalisable to other types.
return (self._lhs_expr * 0) + weight.sort_by(self._lhs_expr)

df = pl.DataFrame(
data=dict(
weight=[5.1, 4.1, 3.1, 2.1, 1.1], a=[1, 5, 3, 4, 2], b=[-3, -2, -5, -2, -4]
)
)

# Working.
print(
"Following works.  Note that the reordering of weight is labeled as 'a' as required."
)
print(
df.select(
pl.col("weight"),
pl.col("a").new.reorder_1(pl.col("weight")),
)
)

# First failure mode, despite the .alias() inside .reorder_1(), the output name of .new.reorder_1() is weight_reorder
print(
"Following almost works, except the output name is not a_reorder, but weight_reorder!"
)
print(
df.select(
pl.col("a"),
pl.col("weight"),
pl.col("a").new.reorder_1(pl.col("weight")).name.suffix("_reorder"),
)
)

# Second (exception) failure mode, using a selector means the names don't get correctly fixed.
print(
"The following throws an exception, as it cannot determine the names of all the columns."
)
try:
print(df.select(pl.selectors.numeric().new.reorder_1(pl.col("weight"))))
except pl.exceptions.ComputeError as e:
print(traceback.format_exc())

# Working kludge.
print(
"The following 'works', including the name flowing through to .suffix() but it cannot possibly be the right way to do this."
)
print(df.select(pl.selectors.numeric().new.reorder_2(pl.col("weight")).name.suffix("_")))

Эта проблема также очевидна в самой документации по Polars, в примере декоратора Register_expr_namespace, воспроизведенном ниже.

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

def register_expr_namespace(name: str) -> Callable[[type[NS]], type[NS]]:
"""
Decorator for registering custom functionality with a Polars Expr.

Parameters
----------
name
Name under which the functionality will be accessed.

See Also
--------
register_dataframe_namespace: Register functionality on a DataFrame.
register_lazyframe_namespace: Register functionality on a LazyFrame.
register_series_namespace: Register functionality on a Series.

Examples
--------
>>> @pl.api.register_expr_namespace("pow_n")
... class PowersOfN:
...     def __init__(self, expr: pl.Expr):
...         self._expr = expr
...
...     def next(self, p: int) -> pl.Expr:
...         return (p ** (self._expr.log(p).ceil()).cast(pl.Int64)).cast(pl.Int64)
...
...     def previous(self, p: int) -> pl.Expr:
...         return (p ** (self._expr.log(p).floor()).cast(pl.Int64)).cast(pl.Int64)
...
...     def nearest(self, p: int) -> pl.Expr:
...          return (p ** (self._expr.log(p)).round(0).cast(pl.Int64)).cast(pl.Int64)
>>>
>>> df = pl.DataFrame([1.4, 24.3, 55.0, 64.001], schema=["n"])
>>> df.select(
...     pl.col("n"),
...     pl.col("n").pow_n.next(p=2).alias("next_pow2"),
...     pl.col("n").pow_n.previous(p=2).alias("prev_pow2"),
...     pl.col("n").pow_n.nearest(p=2).alias("nearest_pow2"),
... )
shape: (4, 4)
┌────────┬───────────┬───────────┬──────────────┐
│ n      ┆ next_pow2 ┆ prev_pow2 ┆ nearest_pow2 │
│ ---    ┆ ---       ┆ ---       ┆ ---          │
│ f64    ┆ i64       ┆ i64       ┆ i64          │
╞════════╪═══════════╪═══════════╪══════════════╡
│ 1.4    ┆ 2         ┆ 1         ┆ 1            │
│ 24.3   ┆ 32        ┆ 16        ┆ 32           │
│ 55.0   ┆ 64        ┆ 32        ┆ 64           │
│ 64.001 ┆ 128       ┆ 64        ┆ 64           │
└────────┴───────────┴───────────┴──────────────┘

В приведенном выше примере, поскольку имена результирующих столбцов явно заданы, все работает нормально, однако имена по умолчанию являются «буквальными», а не более интуитивно понятным именем, каким бы ни было self._expr.
Это до безумия неприятная проблема, и на самом деле неясно, каким должно быть правильное решение.>

Подробнее здесь: https://stackoverflow.com/questions/790 ... -in-polars