PyCall.rb + Pandas: DataFrame#query の代わりにS式っぽく書けないか試してみた

:::note
（追記） @heronshoes さんが [RedAmber](https://github.com/red-data-tools/red_amber) 版を書いてくださいました！ あわせてお読みいただければと思います。

[RedAmberで書いたらこうなる - pandasのquery - Qiita](https://qiita.com/heronshoes/items/962e312146dac8673b49)
:::




# 経緯

これは [データサイエンス100本ノック（構造化データ加工編）](https://github.com/The-Japan-DataScientist-Society/100knocks-preprocess) 例題 P-005 Python 版の解答例です。

```python
df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
            .query('customer_id == "CS018205000001" & amount >= 1000')
```

PyCall.rb + Pandas でもほぼそのまま動くかな……と思いきや、以下を実行するとエラーになります。

```rb
df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']]
    .query('customer_id == "CS018205000001" & amount >= 1000')
#=> PyCall::PyError: <class 'ValueError'>: call stack is not deep enough
```

<details><summary>スタックトレース</summary>

```
PyCall::PyError: <class 'ValueError'>: call stack is not deep enough
  File "/opt/conda/lib/python3.10/site-packages/pandas/util/_decorators.py", line 331, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pandas/core/frame.py", line 4474, in query
    res = self.eval(expr, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pandas/util/_decorators.py", line 331, in wrapper
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pandas/core/frame.py", line 4612, in eval
    return _eval(expr, inplace=inplace, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pandas/core/computation/eval.py", line 345, in eval
    env = ensure_scope(
  File "/opt/conda/lib/python3.10/site-packages/pandas/core/computation/scope.py", line 25, in ensure_scope
    return Scope(
  File "/opt/conda/lib/python3.10/site-packages/pandas/core/computation/scope.py", line 131, in __init__
    frame = sys._getframe(self.level)

/home/jovyan/.local/share/gem/ruby/3.1.0/gems/pycall-1.4.2/lib/pycall/pyobject_wrapper.rb:49:in `query'
/home/jovyan/.local/share/gem/ruby/3.1.0/gems/pycall-1.4.2/lib/pycall/pyobject_wrapper.rb:49:in `method_missing'
(irb):3:in `<top (required)>'
```
</details>

（追記） pandas.rb の issue で報告済みです（[Issue #24](https://github.com/mrkn/pandas.rb/issues/24)）。mrkn さん、お声がけありがとうございました。



---

次のように書くことで同等の操作ができるようです。

```rb
df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
    [(df_receipt['customer_id'] == "CS018205000001") & (df_receipt['amount'] >= 1000)]
```

できますが、何度も `df_receipt['...']` と書かないといけないのが煩雑ですし、DRY でないのが気になりますね。

短い変数名を使えば多少は緩和されますが、DRY でないのは同じですし、わざわざ別の変数を用意したくない場合もあるでしょう。

```rb
df = df_receipt
df[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
    [(df['customer_id'] == "CS018205000001") & (df['amount'] >= 1000)]
```




# S式っぽく書けるようにする

query メソッドに渡す文字列を自前で解析して DataFrame#query と同じ動作を実現する方法も考えられますが、ちょっと大変そうですね。

簡易な対処として、次のようにS式っぽく書けないでしょうか？

```ruby
df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']]
    ._query(
        [:&,
            [:==, :customer_id, "CS018205000001"],
            [:>=, :amount, 1000]]
    )
```

というわけで試しに書いてみたのが以下です。

```ruby
class Pandas::DataFrame
    def _eval_array(xs)
        case xs
        in [:==, lhs, rhs] then _eval(lhs) == _eval(rhs)
        in [:!=, lhs, rhs] then _eval(lhs) != _eval(rhs)
        in [:>=, lhs, rhs] then _eval(lhs) >= _eval(rhs)
        in [:<=, lhs, rhs] then _eval(lhs) <= _eval(rhs)
        in [:> , lhs, rhs] then _eval(lhs) >  _eval(rhs)
        in [:< , lhs, rhs] then _eval(lhs) <  _eval(rhs)
        in [:&, head, *tail]
            tail.inject(_eval(head)){ |result, expr| result & _eval(expr) }
        in [:|, head, *tail]
            tail.inject(_eval(head)){ |result, expr| result | _eval(expr) }
        else
            raise "unsupported (#{xs.inspect})"
        end
    end
    
    def _eval_symbol(sym)
        if self.columns.to_a.include?(sym.to_s)
            self[sym]
        else
            raise "no such column (#{sym})"
        end
    end
    
    def _eval(expr)
        case expr
        when String, Integer then expr
        when Array then _eval_array(expr)
        when Symbol then _eval_symbol(expr)
        else
            raise "unsupported (#{expr.class}) (#{expr.inspect})"
        end
    end
    
    def _query(expr)
        series = _eval(expr)
        self[series]
    end
end
```

軽く使った範囲ではうまく動いているようです。

---

シンボルは列名とみなす、という方式にしてみました。

たとえば `[:==, :customer_id, "CS018205000001"]` の場合、 `:customer_id` は列名とみなされますが、 `"CS018205000001"` は文字列なので列名とはみなされません。

---

式の中で変数を参照できます。

```ruby
amount_min = 1000

df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']]
    ._query(
        [:&,
            [:==, :customer_id, "CS018205000001"],
            [:>=, :amount, amount_min]]
    )
```





# できないこと

## Python のメソッドが使われているパターン

```python
# データサイエンス100本ノック（構造化データ加工編） 例題 P-011 Python 版の解答例
df_customer.query("customer_id.str.endswith('1')", engine='python').head(10)
```

こういうのは対応できてません。今回はパス。




## 左辺が Series ではないパターン

```ruby
df_receipt._query(
    [:<=, 1000, :amount]
)
#=> NoMethodError: undefined method `<=' for #<PyCall::PyObjectWrapper::SwappedOperationAdapter:0x00007f182002c800 @obj=1000>
```

左辺はおそらく `pandas.core.series.Series` じゃないといけないので ↑こういうのはNGです。

```ruby
df_receipt._query(
    [:>=, :amount, 1000]
)
```

↑ こうすればOK。

以下のように `_eval_array` メソッド側で対応することも可能そうです。

```ruby
    def _eval_array(xs)
        case xs
        # ...
        in [:<=, lhs, rhs]
            # 雑な対応。もっとよい方法があると思います。
            begin
                _eval(lhs) <= _eval(rhs)
            rescue
                _eval(rhs) >= _eval(lhs)
            end
        # ...
        end
    end
```





# メモ

## 雑多

- 今回は PyCall.rb + Pandas で試したが、他のデータフレームライブラリでも応用できそう？
- 今回扱った範囲程度なら中置記法にしてもよさそう？


## df['列名'] などの挙動

```ruby
s1 = df_receipt['amount']
p s1.class #=> <class 'pandas.core.series.Series'>

s2 = s1.>(1000)
p s2.class #=> <class 'pandas.core.series.Series'>

df = df_receipt[s2]
p df.class #=> <class 'pandas.core.frame.DataFrame'>
```

なるほど。これをまとめて書いたのが以下。

```ruby
df_receipt[ df_receipt['amount'] > 1000 ]
```



## キーワード引数 level を指定する方法

```ruby
df_receipt[['sales_ymd', 'customer_id', 'product_cd', 'amount']]
    .query('customer_id == "CS018205000001" & amount >= 1000', level: -1)
```

こうするとエラーは発生しなくなるのですが……

- Python, Pandas に詳しくないので妥当な対処方法なのか判断できない
- `@hoge` で変数を参照できない



# 別の簡易な方法 1（2023-05-13 追記）

短い変数名を使いたいが一時的な変数で環境を汚したくないという場合、以下も手軽でよいかもしれませんね。

```rb
def with_temp_df(df)
    yield df
end

with_temp_df(df_receipt) { |df|
    df[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
        [df['customer_id'] == 'CS018205000001']
}
```

あるいは DataFrame クラスのインスタンスメソッドにしてしまうとか。

```rb
class Pandas::DataFrame
    def with_temp_df
        yield self
    end
end

df_receipt.with_temp_df { |df|
    df[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
        [df['customer_id'] == 'CS018205000001']
}
```




# 別の簡易な方法 2（2023-06-04 追記）

@heronshoes さんにコメントで教えていただきました。ありがとうございます。 :pray:

「簡易な方法 1」のように自分でメソッドを定義しなくても `Object#then` （`Object#yield_self`）を使えば同等のことができます。便利！

```rb
df_receipt.then { |df|
    df[['sales_ymd', 'customer_id', 'product_cd', 'amount']] \
        [df['customer_id'] == 'CS018205000001']
}
```

参考: [Object#then (Ruby 3.2 リファレンスマニュアル)](https://docs.ruby-lang.org/ja/3.2/method/Object/i/then.html)


`Object#then` の存在は知っていたのですが、普段使っていないとこういうときにスッと出てこないですね。






# バージョン

```ruby
RUBY_VERSION #=> 3.1.3
PyCall::VERSION #=> 1.4.2
Pandas::VERSION #=> 0.3.8
IRuby::VERSION #=> 0.7.4

sys.version_info #=> sys.version_info(major=3, minor=10, micro=6, releaselevel='final', serial=0)
Pandas.__version__ #=> 1.5.3
```