1つのdataframeにあるの2つの変数の近似値の見つけ方

Question

実現したいこと

200万程度のレコードがあるdf1の変数targetの数値の中で、変数tempの数値に近いレコードを見つけたいのですが、下の配列を使ったソースコードだと結果が返ってくるまでの時間がかかり過ぎるため、より短時間で結果を得られるようにしたいです。
改善案がありましたら教えてください。

該当のソースコード

import numpy as np

def find_closest(arr, val):
    idx = np.abs(arr - val).argmin()
    
    return arr[idx]

    
target = np.array(df1.select('val1').collect())
df2 = df1.withColumn('closest', find_closest(target, temp))

Comments

[melian]

確認ですが、select や collect, withColumn というメソッドを使われているので Polars でしょうか？

[mom_mom]

PySparkです。

[melian]

ああ、すっかり勘違いをしていました。失礼しました。

Answer 1

Polars を利用していると仮定。

import polars as pl
import numpy as np

M = 20_000_000
N = 2_000_000
df1 = pl.DataFrame({'val1': np.random.randint(0, M, N)})
temp = np.random.randint(0, M)

idx = df1.select((pl.col('val1') - temp).abs().arg_min())[0, 0]
df2 = df1[idx]
print(df2)

# shape: (1, 1)
# ┌─────┐
# │ val1     │
# │ ---      │
# │ i64      │
# ╞═════╡
# │ 13116957 │
# └─────┘

Comments

[mom_mom]

ありがとうございます。
Polarsは使っていなかったのですが、Polarsを使うというのも一案ですね。