[GitHub] [beam] TheNeuralBit commented on a change in pull request #13853: [BEAM-9547] Produce better errors for some groupby() and set_index() configurations

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

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



##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -182,14 +194,53 @@ def test_groupby_project(self):
     self._run_test(lambda df: df.groupby('group')['foo'].median(), df)
     self._run_test(lambda df: df.groupby('group')['baz'].median(), df)
     self._run_test(lambda df: df.groupby('group')[['bar', 'baz']].median(), df)
+
+  def test_groupby_errors(self):
+    df = pd.DataFrame({
+        'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)],
+        'foo': [None if i % 11 == 0 else i for i in range(100)],
+        'bar': [None if i % 7 == 0 else 99 - i for i in range(100)],
+        'baz': [None if i % 13 == 0 else i * 2 for i in range(100)],
+    })
+
+    # non-existent projection column
     self._run_test(
         lambda df: df.groupby('group')[['bar', 'baz']].bar.median(),
         df,
         expect_error=True)
     self._run_test(
-        lambda df: df.groupby('group')[['bat']].median(), df, expect_error=True)
+        lambda df: df.groupby('group')[['bad']].median(), df, expect_error=True)
+
+    self._run_test(
+        lambda df: df.groupby('group').bad.median(), df, expect_error=True)
+
+    # non-existent grouping label

Review comment:
       Done

##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -63,23 +63,35 @@ def _run_test(self, func, *args, distributed=True, expect_error=False):
             "Expected an error:\n{expected}\nbut successfully "
             f"returned:\n{actual}")
 
-    if not expect_error:
-      if hasattr(expected, 'equals'):
+    if expect_error:
+      if not isinstance(actual,
+                        type(expected)) or not str(actual) == str(expected):
+        raise AssertionError(
+            f'Expected {expected!r} to be raised, but got {actual!r}'
+        ) from actual
+    else:
+      if isinstance(expected, pd.core.generic.NDFrame):
         if distributed:
-          cmp = lambda df: expected.sort_index().equals(df.sort_index())
+          expected = expected.sort_index()
+          actual = actual.sort_index()
+
+        if isinstance(expected, pd.Series):
+          pd.testing.assert_series_equal(expected, actual)
+        elif isinstance(expected, pd.DataFrame):
+          pd.testing.assert_frame_equal(expected, actual)
         else:
-          cmp = expected.equals
-      elif isinstance(expected, float):
-        cmp = lambda x: (math.isnan(x) and math.isnan(expected)
-                         ) or x == expected == 0 or abs(expected - x) / (
-                             abs(expected) + abs(x)) < 1e-8
+          raise ValueError(
+              "Expected value is an NDFrame, but not a Series or DataFrame.")

Review comment:
       Done

##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -182,14 +194,53 @@ def test_groupby_project(self):
     self._run_test(lambda df: df.groupby('group')['foo'].median(), df)
     self._run_test(lambda df: df.groupby('group')['baz'].median(), df)
     self._run_test(lambda df: df.groupby('group')[['bar', 'baz']].median(), df)
+
+  def test_groupby_errors(self):
+    df = pd.DataFrame({
+        'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)],
+        'foo': [None if i % 11 == 0 else i for i in range(100)],
+        'bar': [None if i % 7 == 0 else 99 - i for i in range(100)],
+        'baz': [None if i % 13 == 0 else i * 2 for i in range(100)],
+    })
+
+    # non-existent projection column
     self._run_test(
         lambda df: df.groupby('group')[['bar', 'baz']].bar.median(),
         df,
         expect_error=True)
     self._run_test(
-        lambda df: df.groupby('group')[['bat']].median(), df, expect_error=True)
+        lambda df: df.groupby('group')[['bad']].median(), df, expect_error=True)
+
+    self._run_test(
+        lambda df: df.groupby('group').bad.median(), df, expect_error=True)
+
+    # non-existent grouping label
+    self._run_test(
+        lambda df: df.groupby(['really_bad', 'foo', 'bad']).foo.sum(),
+        df,
+        expect_error=True)
+    self._run_test(
+        lambda df: df.groupby('bad').foo.sum(), df, expect_error=True)
+
+  def test_set_index(self):
+    df = pd.DataFrame({
+        # Generate some unique columns to use for indexes

Review comment:
       We had some tests earlier that missed some partitioning bugs because they used too small inputs. But you're right I probably went overboard with 100 elements. Dropped this down to 20 and made it deterministic.

##########
File path: sdks/python/apache_beam/dataframe/frames_test.py
##########
@@ -332,6 +383,8 @@ def test_dataframe_getitem(self):
 
   def test_loc(self):
     dates = pd.date_range('1/1/2000', periods=8)
+    # We do not preserve the freq attribute on a DateTime index. Is this a bug?

Review comment:
       Done, thanks




----------------------------------------------------------------
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