TheNeuralBit commented on a change in pull request #16256:
URL: https://github.com/apache/beam/pull/16256#discussion_r778476199
##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -4073,6 +4077,10 @@ def fn_wrapper(x, *args, **kwargs):
requires_partition_by=partitionings.Index(levels),
preserves_partition_by=partitionings.Index(self._grouping_indexes)))
+ @frame_base.with_docs_from(DataFrameGroupBy)
+ def pipe(self, func, *args, **kwargs):
+ return func(self, *args, **kwargs)
Review comment:
Note func can also be a tuple of `(callable, str)`, where str is the
name of a kwarg to pass `self` to (see the
[docs](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.pipe.html)).
Could you implement and test that logic as well?
##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -1245,6 +1245,37 @@ def test_idxmax(self):
self._run_test(lambda s2: s2.idxmax(), s2)
self._run_test(lambda s2: s2.idxmax(skipna=False), s2)
+ def test_pipe(self):
+ def df_times(df, column, times):
+ df[column] = df[column] * times
+ return df
+
+ def s_times(s, times):
+ return s * times
+
+ df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]}, index=[0, 1, 2])
+ s = pd.Series([1, 2, 3, 4, 5], index=[0, 1, 2, 3, 4])
+ # s2 = pd.Series([6, 7], index=[5, 6])
+
+ func_1 = df_times
+ func_2 = frames.DeferredDataFrame.sum
+ func_3 = frames.DeferredSeries.sum
Review comment:
I'm very surprised that these work. `_run_test` works by first executing
the lambda with normal pandas to compute the expected the result. I'd think
that would fail in the pandas case, since it's trying to call the Beam function.
I think it's just a fluke that they work, since `sum` is implemented with
`_agg_method`:
https://github.com/apache/beam/blob/5ccb00298dcc591e0af7061a5ab0fbdb1196ca8c/sdks/python/apache_beam/dataframe/frames.py#L1957
Which defers to `self.agg`:
https://github.com/apache/beam/blob/5ccb00298dcc591e0af7061a5ab0fbdb1196ca8c/sdks/python/apache_beam/dataframe/frames.py#L131
The other tests you have here (`func_1`, `func_4`) are sufficient in my
opinion, could you just remove the `frames.Deferred*` ones, since they're a
little circular?
##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -1245,6 +1245,37 @@ def test_idxmax(self):
self._run_test(lambda s2: s2.idxmax(), s2)
self._run_test(lambda s2: s2.idxmax(skipna=False), s2)
+ def test_pipe(self):
+ def df_times(df, column, times):
+ df[column] = df[column] * times
+ return df
+
+ def s_times(s, times):
+ return s * times
+
+ df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]}, index=[0, 1, 2])
+ s = pd.Series([1, 2, 3, 4, 5], index=[0, 1, 2, 3, 4])
+ # s2 = pd.Series([6, 7], index=[5, 6])
+
+ func_1 = df_times
+ func_2 = frames.DeferredDataFrame.sum
+ func_3 = frames.DeferredSeries.sum
+ func_4 = s_times
+ # func_5 = frames.DeferredSeries.append
+
+ self._run_inplace_test(lambda df: df.pipe(func_1, 'A', 2), df)
+ self._run_test(lambda df: df.pipe(func_2), df)
+ # type assert fails when axis=1
+ # self._run_test(lambda df: df.pipe(func_2, axis=1), df)
+ self._run_inplace_test(lambda df: df.pipe(func_1, 'A', 2).pipe(func_2), df)
+ self._run_test(lambda df: df.pipe(func_2).pipe(func_3), df)
+
+ self._run_test(lambda s: s.pipe(func_4, 2), s)
+ self._run_test(lambda s: s.pipe(func_3), s)
+ self._run_test(lambda s: s.pipe(func_4, 2).pipe(func_3), s)
+ # Can't append non-deferred series
+ # self._run_test(lambda s: s.pipe(func_5, s2).pipe(func_4, 2), s)
Review comment:
You can pass multiple frames to `_run_test` and they will all be
converted to Deferred counterparts and passed to the lambda (for the
distributed run). I think this should work:
```suggestion
self._run_test(lambda s, s2: s.pipe(func_5, s2).pipe(func_4, 2), s, s2)
```
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
