（追記） @heronshoes さんが RedAmber 版を書いてくださいました！あわせてお読みいただければと思います。

RedAmberで書いたらこうなる - pandasのquery - Qiita

経緯

これはデータサイエンス100本ノック（構造化データ加工編）例題 P-005 Python 版の解答例です。

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

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

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

スタックトレース

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)>'

（追記） pandas.rb の issue で報告済みです（Issue #24）。mrkn さん、お声がけありがとうございました。

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

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

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

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

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

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

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

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

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

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

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" は文字列なので列名とはみなされません。

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

amount_min = 1000

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

できないこと

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

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

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

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

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

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

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

↑ こうすればOK。

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

    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['列名'] などの挙動

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'>

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

df_receipt[ df_receipt['amount'] > 1000 ]

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

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

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

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

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

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

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 クラスのインスタンスメソッドにしてしまうとか。

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 さんにコメントで教えていただきました。ありがとうございます。

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

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

参考: Object#then (Ruby 3.2 リファレンスマニュアル)

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

バージョン

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