1. 正規表現で列選択したい

データフレームを扱うときに正規表現で列選択をするケースは多い。

pysparkだけではなくpandasのデータフレームでも同様のケースはあるけれど、今回はpysparkで実現方法を試してみる。

pysparkの専用のメソッドもあるけれども使い勝手が良くなく、最終的にはreモジュールとリスト内包表記を使用した方法が良さそう。

この方法なら、pandas.DataFrameでの同様の処理が可能なのでそういう意味でも汎用的。

2. データ準備

今回もsklearnのbostonデータセットを使用する。

pandas.DataFrameを経由してpyspark.sql.DataFrameを作成しています。

2.1. データ作成

import re
import pandas as pd
from sklearn.datasets import load_boston

data = load_boston()
boston_pdf = pd.DataFrame(
    data.data,
    columns=data.feature_names
)
boston_sdf = spark.createDataFrame(
    boston_pdf
)

2.2. データスキーマ

ouput
In [3]: boston_sdf.printSchema()
root
 |-- CRIM: double (nullable = true)
 |-- ZN: double (nullable = true)
 |-- INDUS: double (nullable = true)
 |-- CHAS: double (nullable = true)
 |-- NOX: double (nullable = true)
 |-- RM: double (nullable = true)
 |-- AGE: double (nullable = true)
 |-- DIS: double (nullable = true)
 |-- RAD: double (nullable = true)
 |-- TAX: double (nullable = true)
 |-- PTRATIO: double (nullable = true)
 |-- B: double (nullable = true)
 |-- LSTAT: double (nullable = true)

3. 試す列選択

列名の最後がSで終了する列を正規表現を使用して選択してみます。

4. pyspark.sql.DataFrame.colRegex

専用のメソッド’colRegex’があります。

pyspark.sql.DataFrame.colRegex

4.1. 実例

colRegexを試してみます。

正規表現の指定方法に癖があり、バッククォートで囲む必要があります。

この指定方法は直感的ではないしバッククォートの入力も面倒なので使い勝手がよくありません。

pyspark専用のメソッドで使い勝手の悪い方法を常用していくのは避けたいところです。

selected_by_colRegex = (
    boston_sdf
    .select(
        boston_sdf.colRegex("`.*S$`")
    )
)
selected_by_colRegex.printSchema()
ouput
root
 |-- INDUS: double (nullable = true)
 |-- CHAS: double (nullable = true)
 |-- DIS: double (nullable = true)

5. reモジュールとリスト内包表記

colRegexの使い勝手が悪いので別の方法を試します。

reモジュールとリスト内包表記を使用した方法です。

この方法は一般的なモジュールと記法を使用しているのでpyspark限定ではなく汎用的です。

もちろんpandas.DataFrameでも使用可能です。

5.1. 実例

selected = (
    boston_sdf
    .select(
        [
            c for c in
            boston_sdf.columns
            if re.search(
                'S$',
                c,
            )
        ]
    )
)
selected.printSchema()
ouput
root
 |-- INDUS: double (nullable = true)
 |-- CHAS: double (nullable = true)
 |-- DIS: double (nullable = true)

6. まとめ

pyspark専用のcolRegexを試してみましたが使い勝手が良くありません。

今後はreモジュール＋リスト内包表記を常用していきます。