feat(arrays): add repeat method equivalent to __mul__/__rmul__

ibis-project · Jul 31, 2023 · b457c7b · b457c7b
1 parent 0ed0ab1
commit b457c7b
Showing 1 changed file with 5 additions and 31 deletions.
diff --git a/ibis/expr/types/arrays.py b/ibis/expr/types/arrays.py
@@ -207,7 +207,7 @@ def __add__(self, other: ArrayValue) -> ArrayValue:
     def __radd__(self, other: ArrayValue) -> ArrayValue:
         return ops.ArrayConcat((other, self)).to_expr()
 
-    def __mul__(self, n: int | ir.IntegerValue) -> ArrayValue:
+    def repeat(self, n: int | ir.IntegerValue) -> ArrayValue:
         """Repeat this array `n` times.
 
         Parameters
@@ -235,7 +235,7 @@ def __mul__(self, n: int | ir.IntegerValue) -> ArrayValue:
         │ [3]                  │
         │ NULL                 │
         └──────────────────────┘
-        >>> t.a * 2
+        >>> t.a.repeat(2)
         ┏━━━━━━━━━━━━━━━━━━━━━━┓
         ┃ ArrayRepeat(a, 2)    ┃
         ┡━━━━━━━━━━━━━━━━━━━━━━┩
@@ -245,37 +245,9 @@ def __mul__(self, n: int | ir.IntegerValue) -> ArrayValue:
         │ [3, 3]               │
         │ []                   │
         └──────────────────────┘
-        """
-        return ops.ArrayRepeat(self, n).to_expr()
-
-    def __rmul__(self, n: int | ir.IntegerValue) -> ArrayValue:
-        """Repeat this array `n` times.
 
-        Parameters
-        ----------
-        n
-            Number of times to repeat `self`.
+        `repeat` is also available using the `*` operator
 
-        Returns
-        -------
-        ArrayValue
-            `self` repeated `n` times
-
-        Examples
-        --------
-        >>> import ibis
-        >>> ibis.options.interactive = True
-        >>> t = ibis.memtable({"a": [[7], [3] , None]})
-        >>> t
-        ┏━━━━━━━━━━━━━━━━━━━━━━┓
-        ┃ a                    ┃
-        ┡━━━━━━━━━━━━━━━━━━━━━━┩
-        │ array<int64>         │
-        ├──────────────────────┤
-        │ [7]                  │
-        │ [3]                  │
-        │ NULL                 │
-        └──────────────────────┘
         >>> 2 * t.a
         ┏━━━━━━━━━━━━━━━━━━━━━━┓
         ┃ ArrayRepeat(a, 2)    ┃
@@ -289,6 +261,8 @@ def __rmul__(self, n: int | ir.IntegerValue) -> ArrayValue:
         """
         return ops.ArrayRepeat(self, n).to_expr()
 
+    __mul__ = __rmul__ = repeat
+
     def unnest(self) -> ir.Value:
         """Flatten an array into a column.