[GitHub] [beam] robertwb commented on a change in pull request #12990: [BEAM-9547] Implement dot method for dataframes.

[GitBox <gi...@apache.org>]

robertwb commented on a change in pull request #12990:
URL: https://github.com/apache/beam/pull/12990#discussion_r499060394



##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -34,6 +34,45 @@ def __array__(self, dtype=None):
 
   between = frame_base._elementwise_method('between')
 
+  def dot(self, other):
+    left = self._expr
+    if isinstance(other, DeferredSeries):
+      right = expressions.ComputedExpression(
+          'to_dataframe',
+          pd.DataFrame, [other._expr],
+          requires_partition_by=partitionings.Nothing(),
+          preserves_partition_by=partitionings.Index())
+      right_is_series = True
+    elif isinstance(other, DeferredDataFrame):
+      right = other._expr
+      right_is_series = False
+    else:
+      raise frame_base.WontImplementError('non-deferred result')
+
+    dots = expressions.ComputedExpression(
+        'dot',
+        # Transpose so we can sum across rows.
+        lambda left,
+        right: pd.DataFrame(left @ right).T,

Review comment:
       Yes, this is the auto-formatter. I'll play around with this.

##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -34,6 +34,45 @@ def __array__(self, dtype=None):
 
   between = frame_base._elementwise_method('between')
 
+  def dot(self, other):
+    left = self._expr
+    if isinstance(other, DeferredSeries):
+      right = expressions.ComputedExpression(
+          'to_dataframe',
+          pd.DataFrame, [other._expr],
+          requires_partition_by=partitionings.Nothing(),
+          preserves_partition_by=partitionings.Index())
+      right_is_series = True
+    elif isinstance(other, DeferredDataFrame):
+      right = other._expr
+      right_is_series = False
+    else:
+      raise frame_base.WontImplementError('non-deferred result')
+
+    dots = expressions.ComputedExpression(
+        'dot',
+        # Transpose so we can sum across rows.
+        lambda left,
+        right: pd.DataFrame(left @ right).T,
+        [left, right],
+        requires_partition_by=partitionings.Index())
+    with expressions.allow_non_parallel_operations(True):
+      sums = expressions.ComputedExpression(
+          'sum',
+          lambda dots: dots.sum(), [dots],
+          requires_partition_by=partitionings.Singleton())
+
+      if right_is_series:
+        result = expressions.ComputedExpression(
+            'extract',
+            lambda df: df[0], [sums],

Review comment:
       Same.

##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -415,6 +453,37 @@ def aggregate(self, func, axis=0, *args, **kwargs):
       'order-sensitive')
   diff = frame_base.wont_implement_method('order-sensitive')
 
+  def dot(self, other):
+    # We want to broadcast the right hand side to all partitions of the left.
+    # This is OK, as its index must be the same size as the columns set of self,
+    # so cannot be too large.
+    class AsScalar(object):
+      def __init__(self, value):
+        self.value = value
+
+    if isinstance(other, frame_base.DeferredFrame):
+      proxy = other._expr.proxy()
+      with expressions.allow_non_parallel_operations():
+        side = expressions.ComputedExpression(
+            'as_scalar',
+            lambda df: AsScalar(df),
+            [other._expr],
+            requires_partition_by=partitionings.Singleton())
+    else:

Review comment:
       This side is all the constant types that pandas may accept. 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org