時系列のデータから高速にデータを取り出して結合するには（SQL, pandas)

Question

###前提・実現したいこと
以前も同じような質問をこちらでさせて頂いたのですが、不足があったので、再度皆様のお力を貸して頂ければと思います。
以下のような時系列のデータ(X)があった際に、先頭フラグ(Y)を基準として、そこから何サンプル毎、というような取り出し方をして新しいデータテーブルを作る際に、一発で変換する方法はありますでしょうか。SQL, pandas問いません。現在、for文を用いてstackをしていますが、処理が遅く困っています。また、前回cumsumとheadを使用した方法を教えて頂いたのですが、headでは実現したい処理に少し手が届かなく困っております。

###対象のデータ
X = 256, 250, 240, 230, 240, 244, 250, 260, 255, 264, 222, 243, 255
Y = 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0

Y = 1の位置から5つ毎のデータを取得してテーブルにする

###実現したい変換先データ
X = 256, 250, 240, 230, 240, 240, 244, 250, 260, 255, 260, 255, 264, 222, 243
L = 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3

※Lは新データの番号

###stack思考のコード

import pandas as pd

X = [256, 250, 240, 230, 240, 244, 250, 260, 255, 264, 222, 243, 255]
Y = [ 1,   0,   0,   0,   1,   0,   0,   1,   0,   0,   0,   0,   0]

data = pd.DataFrame()
data["X"] = X
data["Y"] = Y

print(data)

newX = []
newL = []
for label, index in enumerate(data[data["Y"]==1].index):
    newX.extend(data["X"][index:index+5].values)
    for i in range(5):
        newL.append(label)

newData = pd.DataFrame()
newData["X"] = newX
newData["L"] = newL
print(newData)

###headを用いたコード

import pandas as pd

df = pd.DataFrame({
    'X': [256, 250, 240, 230, 240, 244, 250, 260, 255, 264, 222, 243, 255],
    'Y': [ 1,   0,   0,   0,   1,   0,   0,   1,   0,   0,  0,  0,  0]})

df['L'] = (df['Y'] == 1).cumsum()

df_new = df[df['L'] > 0].groupby('L').head(5) ##Lのラベルの範囲を超えているので取得できない

Answer 1

SQLでの実装例

■PostgreSQL 9.4 以降

SELECT X
     , (N - 1) / 5 + 1 L  -- 「 5 」というマジックナンバーがパラメータ
     , N  -- この列はデモ用に表示させただけで、実際は射影する必要無し
FROM unnest(
    -- 先頭フラグ(Y)は不要
    -- ( X の要素数に応じた結果セットが自動で返ります )
    
    -- 配列も実際はパラメータになるでしょう
    array[256,250,240,230,240,240,244,250,260,255,260,255,264,222,243]::int[]
    -- 配列は以下のように記述しても可
    -- '{256,250,240,230,240,240,244,250,260,255,260,255,264,222,243}'::int[]
)
with ordinality
AS
    db( X, N )
;

■結果

X	L	N
256	1	1
250	1	2
240	1	3
230	1	4
240	1	5
240	2	6
244	2	7
250	2	8
260	2	9
255	2	10
260	3	11
255	3	12
264	3	13
222	3	14
243	3	15

なお、SQLの実行速度については
バージョン 12 で関数のオプティマイザーサポートが導入されたようなので
投入するデータの件数とバージョン次第で大きな差がありそうです。

また、SQLの実行方法は大きく分けて二通りあるでしょう。

1. pythonで動的にSQL文を組み立ててDBに直接クエリを投げる
2. DB側にユーザー定義のstored functionを登録しておいてpythonから呼び出す

以下、PL/PgSQLで記述したストアドのサンプルです

/*
-- 呼び出し例( 使い方 ):
SELECT get_separate_data(
         5,
         array[256,250,240,230,240,240,244,250,260,255,260,255,264,222,243]
);
*/
CREATE OR REPLACE FUNCTION get_separate_data
                  ( delimiter int, paramarray int[] )
    -- 戻り値は CREATE TYPE で定義しても SETOF record と記述してもいいでしょう
    RETURNS TABLE ( X int, L int )
AS
$$
    SELECT X
         , (N - 1) / $1 + 1
    FROM
        unnest ( $2 )
            with ordinality
        AS
            db( X, N )
$$
LANGUAGE SQL;

Comments

[akirajyunia]

ご回答頂きありがとうございました。ベストアンサーに選ばせて頂きました。
unnestとordinalityを使用すればよかったんですね！　ありがとうございました。

[bsdfan]

こちらの回答には、Yに応じてXを作成する部分が入っていないように見えます。
回答には、変換された後のXが埋め込まれているようです。

Answer 2

取り出したい行のインデックスを作成してから、一回で取り出せば早くなると思います。

インデックスの計算は、そのほうが簡単にかけるためnumpyを使っています。

import pandas as pd
import numpy as np

X = [256, 250, 240, 230, 240, 244, 250, 260, 255, 264, 222, 243, 255]
Y = [ 1,   0,   0,   0,   1,   0,   0,   1,   0,   0,   0,   0,   0]

data = pd.DataFrame()
data["X"] = X
data["Y"] = Y

N = 5

idx0, = np.nonzero(data['Y'].array == 1)
idx = np.ravel(np.arange(N) + idx0[:, None])

newData = pd.DataFrame({
    'X': data['X'].array[idx],
    'L': np.arange(len(idx)) // N + 1})

idx0 は numpy.nonzero() で y=1 のインデックス[0, 4, ...]をとりだしていて、
idx は numpy のブロードキャストを使って [0, 1, 2, 3, 4, 4 + 0, 4 + 1, ...] を作っています。

■追記

idxがXの範囲を超える(Y=1のあと5行ない)場合があるケースでは、最後を下記に変更ください。

mask = idx < len(data['X'])
newData = pd.DataFrame({
    'X': data['X'].array[idx[mask]],
    'L': np.arange(len(idx))[mask] // N + 1})

Comments

[akirajyunia]

すぐにご回答頂きまして、ありがとうございました。とても助かりました。
重複したindexを先に作ってしまえばよかったのですね。
SQLの方の回答とベストアンサーで大変悩みましたが、SQLの方が一発で出来るので、そちらをベストアンサーにさせて頂きました。高評価は押させて頂きました。ありがとうございました。

Answer 3

あまり綺麗ではありませんが、一応動いています。

>>> import pandas as pd
>>> X = 256, 250, 240, 230, 240, 244, 250, 260, 255, 264, 222, 243, 255
>>> Y =   1,   0,   0,   0,   1,   0,   0,   1,   0,   0,   0,   0,   0
>>> N = 5
>>>
>>> df = pd.DataFrame({'X': X, 'Y': Y})
>>> df2 = pd.concat([df]+[df['X'].shift(-1-i) for i in range(N-1)], axis=1)
>>> df3 = df2[df2['Y'] == 1].dropna().astype(int)
>>> df3['L'] = df3['Y'].cumsum()
>>> df4 = df3.drop('Y', axis=1).set_index('L').stack().reset_index().drop('level_1', axis=1).rename(columns={0:'X'})
>>> print(df4)
    L    X
0   1  256
1   1  250
2   1  240
3   1  230
4   1  240
5   2  240
6   2  244
7   2  250
8   2  260
9   2  255
10  3  260
11  3  255
12  3  264
13  3  222
14  3  243

Comments

[akirajyunia]

早急なご回答ありがとうございました。df2を出力するところで、concatとfor文を内包表記で書いているのは、とても勉強になりました。ありがとうございました。