pd.DataFrame.__getitem__(self, key)

Nov 28, 2020 15:04 · 484 words · 3 minute read

引数・返り値の型が多様なので、自分のために整理しておく。

前提知識

Python では list や dict や str の要素へのアクセス時によく使われる鉤括弧を使った記法であるが、

l = [9, 4, 0, 7]
l[2] # 0

d = { 'a': 3, 'b': 8 }
d['b'] # 8

s = 'abcde'
s[3] # 'd'

これは __getitem__(self, key) で定義できる。

class X(object):
  def __getitem__(self, key):
    return 'abc'

x = X()
x['a'] # 'abc'

DataFrame での実装

https://github.com/pandas-dev/pandas/blob/1fc5efd76ea1f85979fa2364a292d85482e338aa/pandas/core/frame.py#L2990-L3050

公式ドキュメント

https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html

以下、サンプルコードでは df を次のような DataFrame とする。

import pandas as pd
import numpy as np

df = pd.DataFrame({
  'a': [1, 2, 3, 4, 5],
  'b': [3, 2, 5, 1, 4],
})
"""
   a  b
0  1  3
1  2  2
2  3  5
3  4  1
4  5  4
"""

DataFrane[列名のリスト] -> DataFrame

df[列名のリスト] は DataFrame を返す。SQL だと SELECT 列名1, 列名2, ... FROM df に対応するクエリ。index、行数は元の DataFrame と変わらない。

df[['a']] # DataFrame
"""
   a
0  1
1  2
2  3
3  4
4  5
Name: a, dtype: int64
"""

DataFrame[列名] -> Sereis

df[列名] は Series を返す。SQL だと SELECT 列名 FROM df に対応するクエリ。index、行数は元の DataFrame と変わらない。

df['a'] # Series
"""
0    1
1    2
2    3
3    4
4    5
Name: a, dtype: int64
"""

DataFrame[Series] -> DataFrame

真偽値の Series を渡すと、DataFrame を返す。行数はもとの DataFrame 以下となる。

s = pd.Series([True, False, False, True, True])
df[s]
"""
   a  b
0  1  3
3  4  1
4  5  4
"""

この場合 SQL だと SELECT * FROM df JOIN s ON index WHERE s.bool = True に近い雰囲気。

index がマッチしない場合は警告が出る。 UserWarning: Boolean Series key will be reindexed to match DataFrame index.

s = pd.Series([True, False, False, True, True], index=pd.Index([4, 3, 2, 1, 0]))
df[s]
"""
   a  b
0  1  3
1  2  2
4  5  4
"""

Series の長さが DataFrame の行数より短い場合は例外が投げられる。IndexingError: Unalignable boolean Series provided as indexer (index of the boolean Series and of the indexed object do not match).

Series の長さが DataFrame の行数より多い場合は例外は投げられないが、警告が出る。UserWarning: Boolean Series key will be reindexed to match DataFrame index.

よく見る次のような式は、このタイプ。

# 例: a > b の行を抽出する
df[df['a'] > df['b']] # DataFrame
"""
   a  b
3  4  1
4  5  4
"""

df['a'] > df['b'] の型は Series で index は df と同じ。値はすべて bool である。

DataFrame[List<bool>]

DataFrame[Series<bool>] と同じ。

df[[True, False, True, False, True]]
"""
   a  b
0  1  3
2  3  5
4  5  4
"""

DataFrame[DataFrame<bool>]

bool の DataFrame でフィルタリングできる。True 以外の場所は NaN になる。

dfb = pd.DataFrame({
    'a': [True, False, True, False, True],
    'b': [True, True, False, False, True],
})
df[dfb]
"""
     a    b
0  1.0  3.0
1  NaN  2.0
2  3.0  NaN
3  NaN  NaN
4  5.0  4.0
"""
dfb = pd.DataFrame({})
df[dfb]
"""
    a   b
0 NaN NaN
1 NaN NaN
2 NaN NaN
3 NaN NaN
4 NaN NaN
"""

DataFrame[np.ndarray]

df[np.array(['b'])] # df[['b']] と同じ。DataFrame。
"""
   b
0  3
1  2
2  5
3  1
4  4
"""

ただし列名の 0次元の ndarray の場合は Series を返す。

df[np.array('b')] # df['b'] と同じ。Series
"""
0    3
1    2
2    5
3    1
4    4
Name: b, dtype: int64
"""

DataFrame[callable]

callable を実行し、その返り値を key として使う。

def f(d: DataFrame):
  print(len(d))
  return 'b'

df[f] # 関数 f が実行され、df['b'] と同じ結果を返す