Skip to content

Commit 04a554c

Browse files
majiayu000claude
andauthored
BUG: Add logical operators to pd.col Expression class (#63330)
Co-authored-by: Claude Opus 4.5 <[email protected]>
1 parent 1fd184d commit 04a554c

File tree

2 files changed

+90
-0
lines changed

2 files changed

+90
-0
lines changed

pandas/core/col.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,12 @@
3737
"__lt__": "<",
3838
"__eq__": "==",
3939
"__ne__": "!=",
40+
"__and__": "&",
41+
"__rand__": "&",
42+
"__or__": "|",
43+
"__ror__": "|",
44+
"__xor__": "^",
45+
"__rxor__": "^",
4046
}
4147

4248

@@ -157,6 +163,28 @@ def __mod__(self, other: Any) -> Expression:
157163
def __rmod__(self, other: Any) -> Expression:
158164
return self._with_binary_op("__rmod__", other)
159165

166+
# Logical ops
167+
def __and__(self, other: Any) -> Expression:
168+
return self._with_binary_op("__and__", other)
169+
170+
def __rand__(self, other: Any) -> Expression:
171+
return self._with_binary_op("__rand__", other)
172+
173+
def __or__(self, other: Any) -> Expression:
174+
return self._with_binary_op("__or__", other)
175+
176+
def __ror__(self, other: Any) -> Expression:
177+
return self._with_binary_op("__ror__", other)
178+
179+
def __xor__(self, other: Any) -> Expression:
180+
return self._with_binary_op("__xor__", other)
181+
182+
def __rxor__(self, other: Any) -> Expression:
183+
return self._with_binary_op("__rxor__", other)
184+
185+
def __invert__(self) -> Expression:
186+
return Expression(lambda df: ~self(df), f"(~{self._repr_str})")
187+
160188
def __array_ufunc__(
161189
self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any
162190
) -> Expression:

pandas/tests/test_col.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,65 @@ def mean(self):
9797
result = df.assign(b=pd.col("a").xyz.mean())
9898
expected = pd.DataFrame({"a": [1, 2, 3], "b": [2.0, 2.0, 2.0]})
9999
tm.assert_frame_equal(result, expected)
100+
101+
102+
@pytest.mark.parametrize(
103+
("expr", "expected_values", "expected_str"),
104+
[
105+
(
106+
pd.col("a") & pd.col("b"),
107+
[False, False, True, False],
108+
"(col('a') & col('b'))",
109+
),
110+
(
111+
pd.col("a") & True,
112+
[True, False, True, False],
113+
"(col('a') & True)",
114+
),
115+
(
116+
pd.col("a") | pd.col("b"),
117+
[True, True, True, True],
118+
"(col('a') | col('b'))",
119+
),
120+
(
121+
pd.col("a") | False,
122+
[True, False, True, False],
123+
"(col('a') | False)",
124+
),
125+
(
126+
pd.col("a") ^ pd.col("b"),
127+
[True, True, False, True],
128+
"(col('a') ^ col('b'))",
129+
),
130+
(
131+
pd.col("a") ^ True,
132+
[False, True, False, True],
133+
"(col('a') ^ True)",
134+
),
135+
(
136+
~pd.col("a"),
137+
[False, True, False, True],
138+
"(~col('a'))",
139+
),
140+
],
141+
)
142+
def test_col_logical_ops(
143+
expr: Expression, expected_values: list[bool], expected_str: str
144+
) -> None:
145+
# https://github.com/pandas-dev/pandas/issues/63322
146+
df = pd.DataFrame({"a": [True, False, True, False], "b": [False, True, True, True]})
147+
result = df.assign(c=expr)
148+
expected = pd.DataFrame(
149+
{
150+
"a": [True, False, True, False],
151+
"b": [False, True, True, True],
152+
"c": expected_values,
153+
}
154+
)
155+
tm.assert_frame_equal(result, expected)
156+
assert str(expr) == expected_str
157+
158+
# Test that the expression works with .loc
159+
result = df.loc[expr]
160+
expected = df[expected_values]
161+
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)