optimize Series.dot for dataframe with single level columns

Author: shobsi

bigframes/series.py

@@ -680,8 +680,25 @@ def rdivmod(self, other) -> Tuple[Series, Series]:  # type: ignore
         return (self.rfloordiv(other), self.rmod(other))
 
     def __matmul__(self, other):
-        if isinstance(other, bigframes.dataframe.DataFrame):
-            return Series(
+        if isinstance(other, Series):
+            return (self * other).sum()
+
+        # At this point other must be a DataFrame
+        if len(other.columns.names) == 1:
+            # Single level columns in other
+            na_df = other.isna().any()
+            mul_df = Series(
+                [(self * other[col]).sum() for col in other.columns],
+                index=other.columns,
+                name=self.name,
+            )
+            result = mul_df.mask(na_df)
+        else:
+            # Multi level columns in other
+            # TODO(b/313747368): Remove this once DataFrame.any() honors
+            # multi-level index, as the logic in the if clause should generalize
+            # for multi-level columns in other
+            result = Series(
                 [
                     pandas.NA if other[col].isna().any() else (self * other[col]).sum()
                     for col in other.columns
@@ -690,8 +707,7 @@ def __matmul__(self, other):
                 name=self.name,
             )
 
-        # At this point other must be a Series
-        return (self * other).sum()
+        return result
 
     dot = __matmul__
 

tests/system/small/test_multiindex.py

@@ -1055,10 +1055,15 @@ def test_series_dot_df_column_multi_index():
         [["col0", "col0", "col1"], ["col00", "col01", "col11"]]
     )
 
-    bf_result = bpd.Series(left) @ bpd.DataFrame(right, columns=multi_level_columns)
-    pd_result = pandas.Series(left) @ pandas.DataFrame(
-        right, columns=multi_level_columns
-    )
+    bf_left_s = bpd.Series(left)
+    bf_right_df = bpd.DataFrame(right)
+    bf_right_df.columns = multi_level_columns
+    bf_result = bf_left_s @ bf_right_df
+
+    pd_left_s = pandas.Series(left)
+    pd_right_df = pandas.DataFrame(right)
+    pd_right_df.columns = multi_level_columns
+    pd_result = pd_left_s @ pd_right_df
 
     pandas.testing.assert_series_equal(
         bf_result.to_pandas(), pd_result, check_index_type=False, check_dtype=False