Implement dot() for Series

Author: itholic

databricks/koalas/missing/series.py

@@ -60,7 +60,6 @@ class _MissingPandasLikeSeries(object):
     combine_first = unsupported_function('combine_first')
     cov = unsupported_function('cov')
     divmod = unsupported_function('divmod')
-    dot = unsupported_function('dot')
     droplevel = unsupported_function('droplevel')
     duplicated = unsupported_function('duplicated')
     ewm = unsupported_function('ewm')

databricks/koalas/series.py

@@ -4194,6 +4194,70 @@ def pct_change(self, periods=1):
 
         return self._with_new_scol((scol - prev_row) / prev_row)
 
+    def dot(self, other):
+        """
+        Compute the dot product between the Series and the columns of other.
+
+        This method computes the dot product between the Series and another
+        one, or the Series and each columns of a DataFrame.
+
+        It can also be called using `self @ other` in Python >= 3.5.
+
+        Parameters
+        ----------
+        other : Series, DataFrame.
+            The other object to compute the dot product with its columns.
+
+        Returns
+        -------
+        scalar, Series
+            Return the dot product of the Series and other if other is a
+            Series, the Series of the dot product of Series and each rows of
+            other if other is a DataFrame.
+
+        Notes
+        -----
+        The Series and other has to share the same index if other is a Series
+        or a DataFrame.
+
+        Examples
+        --------
+        >>> from databricks.koalas.config import set_option, reset_option
+        >>> set_option("compute.ops_on_diff_frames", True)
+        >>> s = ks.Series([0, 1, 2, 3])
+        >>> other = ks.Series([-1, 2, -3, 4])
+
+        >>> s.dot(other)
+        8
+
+        >>> s @ other
+        8
+
+        >>> df = ks.DataFrame([[0, 1], [-2, 3], [4, -5], [6, 7]])
+        >>> s.dot(df)
+        0    24
+        1    14
+        Name: 0, dtype: int64
+
+        >>> reset_option("compute.ops_on_diff_frames")
+        """
+        if repr(self.index) != repr(other.index):
+            raise ValueError("matrices are not aligned")
+
+        if isinstance(other, DataFrame):
+            idx_val_dict = {col_name: (self * other[col_name]).sum() for col_name in other}
+            result = Series(idx_val_dict)
+        elif isinstance(other, Series):
+            result = (self * other).sum()
+
+        return result
+
+    def __matmul__(self, other):
+        """
+        Matrix multiplication using binary `@` operator in Python>=3.5.
+        """
+        return self.dot(other)
+
     def _cum(self, func, skipna, part_cols=()):
         # This is used to cummin, cummax, cumsum, etc.
         index_columns = self._internal.index_columns

databricks/koalas/tests/test_ops_on_diff_frames.py

@@ -574,6 +574,44 @@ def test_multi_index_column_assignment_frame(self):
         pdf["c"] = 1
         self.assert_eq(repr(kdf), repr(pdf))
 
+    def test_dot(self):
+        kser = ks.Series([90, 91, 85], index=[2, 4, 1])
+        pser = kser.to_pandas()
+        kser_other = ks.Series([90, 91, 85], index=[2, 4, 1])
+        pser_other = kser_other.to_pandas()
+
+        self.assert_eq(kser.dot(kser_other), pser.dot(pser_other))
+
+        kser_other = ks.Series([90, 91, 85], index=[1, 2, 4])
+        with self.assertRaisesRegex(ValueError, "matrices are not aligned"):
+            kser.dot(kser_other)
+
+        kser_other = ks.Series([90, 91, 85, 100], index=[2, 4, 1, 0])
+        with self.assertRaisesRegex(ValueError, "matrices are not aligned"):
+            kser.dot(kser_other)
+
+        # with DataFrame
+        kdf = ks.DataFrame([[0, 1], [-2, 3], [4, -5]], index=[2, 4, 1])
+        pdf = kdf.to_pandas()
+
+        self.assert_eq(kser.dot(kdf), pser.dot(pdf))
+
+        kdf = ks.DataFrame([[0, 1], [-2, 3], [4, -5]], index=[1, 2, 4])
+        with self.assertRaisesRegex(ValueError, "matrices are not aligned"):
+            kser.dot(kdf)
+
+        # for MultiIndex
+        midx = pd.MultiIndex([['lama', 'cow', 'falcon'],
+                              ['speed', 'weight', 'length']],
+                             [[0, 0, 0, 1, 1, 1, 2, 2, 2],
+                              [0, 1, 2, 0, 1, 2, 0, 1, 2]])
+        kser = ks.Series([45, 200, 1.2, 30, 250, 1.5, 320, 1, 0.3], index=midx)
+        pser = kser.to_pandas()
+        kser_other = ks.Series([-450, 20, 12, -30, -250, 15, -320, 100, 3], index=midx)
+        pser_other = kser_other.to_pandas()
+
+        self.assert_eq(kser.dot(kser_other), pser.dot(pser_other))
+
 
 class OpsOnDiffFramesDisabledTest(ReusedSQLTestCase, SQLTestUtils):
 

docs/source/reference/series.rst

@@ -82,6 +82,7 @@ Binary operator functions
    Series.ge
    Series.ne
    Series.eq
+   Series.dot
 
 Function application, GroupBy & Window
 --------------------------------------