[GitHub] [beam] TheNeuralBit commented on a change in pull request #15827: [BEAM-12560] Dataframe idxmin and idxmax implementation

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

TheNeuralBit commented on a change in pull request #15827:
URL: https://github.com/apache/beam/pull/15827#discussion_r751805764



##########
File path: sdks/python/apache_beam/dataframe/frames.py
##########
@@ -1277,6 +1277,90 @@ def align(self, other, join, axis, level, method, **kwargs):
       requires_partition_by=partitionings.Arbitrary(),
       preserves_partition_by=partitionings.Singleton())
 
+  @frame_base.with_docs_from(pd.Series)
+  @frame_base.args_to_kwargs(pd.Series)
+  @frame_base.populate_defaults(pd.Series)
+  def idxmin(self, **kwargs):
+    skipna = kwargs.get('skipna', True)
+
+    def compute_idxmin(s):
+      min_index = s.idxmin(**kwargs)
+      if pd.isna(min_index):
+        return s
+      else:
+        return s.loc[[min_index]]
+
+    # Avoids empty Series error when evaluating proxy
+    index_dtype = self._expr.proxy().index.dtype
+    index = pd.Index([], dtype=index_dtype)
+    proxy = self._expr.proxy().copy()
+    proxy.index = index
+    proxy = proxy.append(
+        pd.Series([np.inf], index=np.asarray(['0']).astype(proxy.index.dtype)))
+
+    if not skipna:
+      proxy = proxy.append(
+          pd.Series([None],
+                    index=np.asarray(['1']).astype(proxy.index.dtype)).astype(
+                        proxy.dtype))
+
+    idx_min = expressions.ComputedExpression(
+        'idx_min',
+        compute_idxmin, [self._expr],
+        proxy=proxy,
+        requires_partition_by=partitionings.Index(),
+        preserves_partition_by=partitionings.Singleton())
+
+    with expressions.allow_non_parallel_operations(True):
+      return frame_base.DeferredFrame.wrap(
+          expressions.ComputedExpression(
+              'idxmin_combine',
+              lambda s: s.idxmin(**kwargs), [idx_min],
+              requires_partition_by=partitionings.Singleton(),
+              preserves_partition_by=partitionings.Singleton()))
+
+  @frame_base.with_docs_from(pd.Series)
+  @frame_base.args_to_kwargs(pd.Series)
+  @frame_base.populate_defaults(pd.Series)
+  def idxmax(self, **kwargs):
+    skipna = kwargs.get('skipna', True)
+
+    def compute_idxmax(s):
+      max_index = s.idxmax(**kwargs)
+      if pd.isna(max_index):
+        return s
+      else:
+        return s.loc[[max_index]]
+
+    # Avoids empty Series error when evaluating proxy
+    index_dtype = self._expr.proxy().index.dtype
+    index = pd.Index([], dtype=index_dtype)
+    proxy = self._expr.proxy().copy()
+    proxy.index = index
+    proxy = proxy.append(
+        pd.Series([-np.inf], index=np.asarray(['0']).astype(proxy.index.dtype)))

Review comment:
       I think you should be able to just propagate the proxy from the input in this case (i.e. `proxy=self._expr.proxy()`), since `compute_idxmax` always returns a Series with the same index and dtype as the input, right?

##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -1100,6 +1100,44 @@ def test_dt_tz_localize_nonexistent(self):
             'Europe/Warsaw', ambiguous='NaT', nonexistent=pd.Timedelta('1H')),
         s)
 
+  def test_idxmin(self):
+    df = pd.DataFrame({
+        'consumption': [10.51, 103.11, 55.48],
+        'co2_emissions': [37.2, 19.66, 1712]
+    },
+                      index=['Pork', 'Wheat Products', 'Beef'])
+
+    s = pd.Series(data=[4, 3, None, 1], index=['A', 'B', 'C', 'D'])
+    s2 = pd.Series(data=[1, 2, 3], index=[1, 2, 3])
+
+    self._run_test(lambda df: df.idxmin(), df, nonparallel=True)

Review comment:
       It looks like you were able to make `Series.idxmin` parallelizable by breaking it up into two steps. Shouldn't you be able to do the same thing for `DataFrame.idxmin` with `axis=0`?




-- 
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: github-unsubscribe@beam.apache.org

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