ENH #60693: shift operations for Series and DataFrames

This commit introduces the rshift and lshift method for both Series and DataFrames in pandas.
It also adds the corresponding in-place methods.
These methdos don't work between a Series and a Dataframe, or if the two Series or DataFrames differ in size.

Co-authored-by: Tomas Macieira <[email protected]>

Author: David-msggc

doc/source/whatsnew/v3.0.0.rst

@@ -79,6 +79,7 @@ Other enhancements
 - Add ``"delete_rows"`` option to ``if_exists`` argument in :meth:`DataFrame.to_sql` deleting all records of the table before inserting data (:issue:`37210`).
 - Added half-year offset classes :class:`HalfYearBegin`, :class:`HalfYearEnd`, :class:`BHalfYearBegin` and :class:`BHalfYearEnd` (:issue:`60928`)
 - Errors occurring during SQL I/O will now throw a generic :class:`.DatabaseError` instead of the raw Exception type from the underlying driver manager library (:issue:`60748`)
+- Implemented :meth:`Series.rshift`, :meth:`Series.lshift`, :meth:`DataFrame.rshift` and :meth:`DataFrame.lshift` for bitwise right and left shift operations (:issue:`60693`)
 - Implemented :meth:`Series.str.isascii` and :meth:`Series.str.isascii` (:issue:`59091`)
 - Improved deprecation message for offset aliases (:issue:`60820`)
 - Multiplying two :class:`DateOffset` objects will now raise a ``TypeError`` instead of a ``RecursionError`` (:issue:`59442`)

pandas/core/arraylike.py

@@ -248,6 +248,22 @@ def __pow__(self, other):
     def __rpow__(self, other):
         return self._arith_method(other, roperator.rpow)
 
+    @unpack_zerodim_and_defer("__rshift__")
+    def __rshift__(self, other):
+        return self._arith_method(other, operator.rshift)
+
+    @unpack_zerodim_and_defer("__rrshift__")
+    def __rrshift__(self, other):
+        return self._arith_method(other, roperator.rrshift)
+
+    @unpack_zerodim_and_defer("__lshift__")
+    def __lshift__(self, other):
+        return self._arith_method(other, operator.lshift)
+
+    @unpack_zerodim_and_defer("__rlshift__")
+    def __rlshift__(self, other):
+        return self._arith_method(other, roperator.rlshift)
+
 
 # -----------------------------------------------------------------------------
 # Helpers to implement __array_ufunc__

pandas/core/computation/expressions.py

@@ -153,6 +153,10 @@ def _evaluate_numexpr(op, op_str, left_op, right_op):
     roperator.rmod: None,
     operator.pow: "**",
     roperator.rpow: "**",
+    operator.rshift: None,
+    roperator.rrshift: None,
+    operator.lshift: None,
+    roperator.rlshift: None,
     operator.eq: "==",
     operator.ne: "!=",
     operator.le: "<=",

pandas/core/frame.py

@@ -8454,6 +8454,38 @@ def rpow(
             other, roperator.rpow, level=level, fill_value=fill_value, axis=axis
         )
 
+    @Appender(ops.make_flex_doc("rshift", "dataframe"))
+    def rshift(
+        self, other, axis: Axis = "columns", level=None, fill_value=None
+    ) -> DataFrame:
+        return self._flex_arith_method(
+            other, operator.rshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("rrshift", "dataframe"))
+    def rrshift(
+        self, other, axis: Axis = "columns", level=None, fill_value=None
+    ) -> DataFrame:
+        return self._flex_arith_method(
+            other, roperator.rrshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("lshift", "dataframe"))
+    def lshift(
+        self, other, axis: Axis = "columns", level=None, fill_value=None
+    ) -> DataFrame:
+        return self._flex_arith_method(
+            other, operator.lshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("rlshift", "dataframe"))
+    def rlshift(
+        self, other, axis: Axis = "columns", level=None, fill_value=None
+    ) -> DataFrame:
+        return self._flex_arith_method(
+            other, roperator.rlshift, level=level, fill_value=fill_value, axis=axis
+        )
+
     # ----------------------------------------------------------------------
     # Combination-Related
 

pandas/core/generic.py

@@ -11705,6 +11705,16 @@ def __ixor__(self, other) -> Self:
         # error: Unsupported left operand type for ^ ("Type[NDFrame]")
         return self._inplace_method(other, type(self).__xor__)  # type: ignore[operator]
 
+    @final
+    def __irshift__(self, other) -> Self:
+        # error: Unsupported left operand type for >> ("Type[NDFrame]")
+        return self._inplace_method(other, type(self).__rshift__)  # type: ignore[operator]
+
+    @final
+    def __ilshift__(self, other) -> Self:
+        # error: Unsupported left operand type for << ("Type[NDFrame]")
+        return self._inplace_method(other, type(self).__lshift__)  # type: ignore[operator]
+
     # ----------------------------------------------------------------------
     # Misc methods
 

pandas/core/ops/array_ops.py

@@ -495,6 +495,8 @@ def get_array_op(op):
         "mod",
         "divmod",
         "pow",
+        "shift",
+        "lshift",
     }:
         return partial(arithmetic_op, op=op)
     else:

pandas/core/ops/docstrings.py

@@ -216,6 +216,31 @@ def make_flex_doc(op_name: str, typ: str) -> str:
 """
 )
 
+_rshift_example_SERIES = (
+    _common_examples_algebra_SERIES
+    + """
+>>> a.rshift(1, fill_value=0)
+a    0.0
+b    0.0
+c    0.0
+d    NaN
+dtype: float64
+"""
+)
+
+_lshift_example_SERIES = (
+    _common_examples_algebra_SERIES
+    + """
+>>> a.rshift(1, fill_value=0)
+a    2.0
+b    2.0
+c    2.0
+d    NaN
+dtype: float64
+"""
+)
+
+
 _ne_example_SERIES = (
     _common_examples_algebra_SERIES
     + """
@@ -365,6 +390,20 @@ def make_flex_doc(op_name: str, typ: str) -> str:
         "series_returns": _returns_tuple,
         "df_examples": None,
     },
+    "rshift": {
+        "op": ">>",
+        "desc": "Bitwise right shift",
+        "reverse": "rrshift",
+        "series_examples": _rshift_example_SERIES,
+        "series_returns": _returns_series,
+    },
+    "lshift": {
+        "op": "<<",
+        "desc": "Bitwise left shift",
+        "reverse": "rlshift",
+        "series_examples": _lshift_example_SERIES,
+        "series_returns": _returns_series,
+    },
     # Comparison Operators
     "eq": {
         "op": "==",

pandas/core/roperator.py

@@ -61,3 +61,11 @@ def ror_(left, right):
 
 def rxor(left, right):
     return operator.xor(right, left)
+
+
+def rrshift(left, right):
+    return operator.rshift(right, left)
+
+
+def rlshift(left, right):
+    return operator.lshift(right, left)

pandas/core/series.py

@@ -6578,6 +6578,30 @@ def rdivmod(self, other, level=None, fill_value=None, axis: Axis = 0) -> Series:
             other, roperator.rdivmod, level=level, fill_value=fill_value, axis=axis
         )
 
+    @Appender(ops.make_flex_doc("rshift", "series"))
+    def rshift(self, other, level=None, fill_value=None, axis: Axis = 0) -> Series:
+        return self._flex_method(
+            other, operator.rshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("rrshift", "series"))
+    def rrshift(self, other, level=None, fill_value=None, axis: Axis = 0) -> Series:
+        return self._flex_method(
+            other, roperator.rrshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("lshift", "series"))
+    def lshift(self, other, level=None, fill_value=None, axis: Axis = 0) -> Series:
+        return self._flex_method(
+            other, operator.lshift, level=level, fill_value=fill_value, axis=axis
+        )
+
+    @Appender(ops.make_flex_doc("rlshift", "series"))
+    def rlshift(self, other, level=None, fill_value=None, axis: Axis = 0) -> Series:
+        return self._flex_method(
+            other, roperator.rlshift, level=level, fill_value=fill_value, axis=axis
+        )
+
     # ----------------------------------------------------------------------
     # Reductions
 

pandas/tests/frame/test_arithmetic.py

@@ -2199,3 +2199,99 @@ def test_mixed_col_index_dtype(using_infer_string):
             dtype = "string"
         expected.columns = expected.columns.astype(dtype)
     tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_rshift():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    result = df >> 2
+    expected = DataFrame({"a": [2, 4], "b": [8, 16]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_rshift_negative():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    result = df >> -1
+    expected = DataFrame({"a": [0, 0], "b": [0, 0]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_rshift_zero():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    result = df >> 0
+
+    tm.assert_frame_equal(result, df)
+
+
+def test_dataframe_lshift():
+    df = DataFrame({"a": [1, 2], "b": [4, 8]})
+    result = df << 2
+    expected = DataFrame({"a": [4, 8], "b": [16, 32]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_lshift_negative():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    result = df << -1
+    expected = DataFrame({"a": [0, 0], "b": [0, 0]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_lshift_zero():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    result = df << 0
+
+    tm.assert_frame_equal(result, df)
+
+
+def test_dataframe_rrshift():
+    df = DataFrame({"a": [1, 2], "b": [3, 4]})
+    result = operator.rshift(64, df)
+    expected = DataFrame({"a": [32, 16], "b": [8, 4]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_rlshift():
+    df = DataFrame({"a": [1, 2], "b": [3, 4]})
+    result = operator.lshift(1, df)
+    expected = DataFrame({"a": [2, 4], "b": [8, 16]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_irshift():
+    df = DataFrame({"a": [8, 16], "b": [32, 64]})
+    df >>= 2
+    expected = DataFrame({"a": [2, 4], "b": [8, 16]})
+
+    tm.assert_frame_equal(df, expected)
+
+
+def test_dataframe_ilshift():
+    df = DataFrame({"a": [1, 2], "b": [4, 8]})
+    df <<= 2
+    expected = DataFrame({"a": [4, 8], "b": [16, 32]})
+
+    tm.assert_frame_equal(df, expected)
+
+
+def test_dataframe_rshift_dataframe():
+    df1 = DataFrame({"a": [8, 4], "b": [2, 1]})
+    df2 = DataFrame({"a": [3, 2], "b": [1, 0]})
+    result = df1 >> df2
+    expected = DataFrame({"a": [1, 1], "b": [1, 1]})
+
+    tm.assert_frame_equal(result, expected)
+
+
+def test_dataframe_lshift_dataframe():
+    df1 = DataFrame({"a": [1, 2], "b": [4, 8]})
+    df2 = DataFrame({"a": [3, 2], "b": [1, 0]})
+    result = df1 << df2
+    expected = DataFrame({"a": [8, 8], "b": [8, 8]})
+
+    tm.assert_frame_equal(result, expected)