从极坐标中的一组列中获取最小元组的有效方法

Question

在极地中最有效的方法是什么：

import polars as pl
import numpy as np
rng = np.random.default_rng()

df = pl.DataFrame([
    pl.Series('a', rng.normal(size=10_000_000)),
    pl.Series('b', rng.normal(size=10_000_000)),
])

df.sort('a', 'b').head(1)

即我想根据字典顺序找到最小的数字元组 (a, b)。求最小值不需要完全排序，所以上面的代码效率很低。我试过了

df.lazy().sort('a', 'b').head(1).collect()

但这并不会带来显着的加速。创建

pl.Struct

类型的单列也没有帮助，因为极坐标似乎没有定义结构排序的概念。

更新

感谢ΩΠΟΚΕΚΡΥΜΜΕΝΟΣ的回答，我现在已经解决了以下解决方案：

def lexicographic_min(df):
    columns = list(df.columns)
    for col in columns:
        if df.select(pl.col(col).is_null().any()).row(0)[0]:
            df = df.filter(pl.col(col).is_null())
        else:
            df = df.filter(pl.col(col) == pl.col(col).min())
    return df.row(0)

def lexicographic_max(df):
    columns = list(df.columns)
    for col in columns:
        df = df.filter(pl.col(col) == pl.col(col).max())
    return df.row(0)

此版本通过将空值视为比任何非空值“小”来处理空值，并且不需要任何调用

.sort()

。

Answer 1

一个简单的加速方法是在连续的步骤中执行过滤。 所以从这个数据开始：

import polars as pl
import numpy as np
import time

rng = np.random.default_rng()
size = 100_000_000
df = pl.DataFrame([
    pl.Series('a', rng.normal(size=size)),
    pl.Series('b', rng.normal(size=size)),
])
df

shape: (100000000, 2)
┌──────────┬───────────┐
│ a        ┆ b         │
│ ---      ┆ ---       │
│ f64      ┆ f64       │
╞══════════╪═══════════╡
│ 0.691714 ┆ 0.042488  │
│ 1.328573 ┆ -0.125364 │
│ -0.23431 ┆ -0.463718 │
│ 0.383145 ┆ -0.307686 │
│ ...      ┆ ...       │
│ 0.12102  ┆ -0.046244 │
│ 0.556771 ┆ 0.208904  │
│ 0.665356 ┆ 1.825813  │
│ 0.660261 ┆ 0.577974  │
└──────────┴───────────┘

这里有两个有趣的时间点。 第一篇来自原帖：

start = time.perf_counter()
df.sort('a', 'b').head(1)
print(time.perf_counter() - start)

>>> df.sort(['a', 'b']).head(1)
shape: (1, 2)
┌───────────┬───────────┐
│ a         ┆ b         │
│ ---       ┆ ---       │
│ f64       ┆ f64       │
╞═══════════╪═══════════╡
│ -5.567192 ┆ -0.658551 │
└───────────┴───────────┘
>>> print(time.perf_counter() - start)
2.884132520000094

与连续过滤...

start = time.perf_counter()
(
    df
    .filter(pl.col('a') == pl.col('a').min())
    .filter(pl.col('b') == pl.col('b').min())
)
print(time.perf_counter() - start)

>>> start = time.perf_counter()
>>> (
...     df
...     .filter(pl.col('a') == pl.col('a').min())
...     .filter(pl.col('b') == pl.col('b').min())
... 
... )
shape: (1, 2)
┌───────────┬───────────┐
│ a         ┆ b         │
│ ---       ┆ ---       │
│ f64       ┆ f64       │
╞═══════════╪═══════════╡
│ -5.567192 ┆ -0.658551 │
└───────────┴───────────┘
>>> print(time.perf_counter() - start)
0.10943432300018685

编辑：空值

这是一种解决方法......但如果您需要将

null

值视为最小值（并且仍想保持性能），我们可以执行以下操作。

假设我们有以下数据：

rng = np.random.default_rng()
size = 100_000_000
null_count = 100_000
df = pl.DataFrame(
    [
        pl.Series("a", rng.normal(size=size))
        .extend_constant(None, null_count)
        .shuffle(seed=1),
        pl.Series("b", rng.normal(size=size))
        .extend_constant(None, null_count)
        .shuffle(seed=2),
    ]
)
df

shape: (100100000, 2)
┌───────────┬───────────┐
│ a         ┆ b         │
│ ---       ┆ ---       │
│ f64       ┆ f64       │
╞═══════════╪═══════════╡
│ -0.336999 ┆ 0.230278  │
│ -0.654095 ┆ -0.056833 │
│ 0.195498  ┆ -1.239257 │
│ 0.490639  ┆ -0.098735 │
│ ...       ┆ ...       │
│ 1.077227  ┆ 1.387681  │
│ -1.75284  ┆ 0.043309  │
│ -0.910276 ┆ -0.762231 │
│ -0.784335 ┆ 0.286225  │
└───────────┴───────────┘

我们可以在过滤步骤之后进行排序。 这应该可以让您保留大部分性能提升（取决于数据中

null

值的比例）。

start = time.perf_counter()
(
    df
    .filter(pl.col('a').is_null() | (pl.col('a') == pl.col('a').min()))
    .sort('a').filter(pl.col('a') == pl.col('a').first())
    .filter(pl.col('b').is_null() | (pl.col('b') == pl.col('b').min()))
    .sort('b').filter(pl.col('b') == pl.col('b').first())
    .head(1)
)
print(time.perf_counter() - start)

shape: (1, 2)
┌──────┬──────┐
│ a    ┆ b    │
│ ---  ┆ ---  │
│ f64  ┆ f64  │
╞══════╪══════╡
│ null ┆ null │
└──────┴──────┘
>>> print(time.perf_counter() - start)
0.1509907710023981

与单独的

sort

进行比较....

>>> start = time.perf_counter()
>>> df.sort('a', 'b').head(1)
shape: (1, 2)
┌──────┬──────┐
│ a    ┆ b    │
│ ---  ┆ ---  │
│ f64  ┆ f64  │
╞══════╪══════╡
│ null ┆ null │
└──────┴──────┘
>>> print(time.perf_counter() - start)
3.319076051997399

这绝对不优雅......但是（取决于您的数据），它可能仍然保留大部分性能提升。