Address further comments.

robertwb · robertwb · commit d76c748b4e2a · 2021-12-30T13:52:18.000-08:00
diff --git a/sdks/python/apache_beam/dataframe/expressions.py b/sdks/python/apache_beam/dataframe/expressions.py
@@ -121,6 +121,7 @@ def evaluate_with(input_partitioning):
         # is computed trivially with Singleton partitioning and passes.
         for input_partitioning in sorted(set([expr.requires_partition_by(),
                                               partitionings.Arbitrary(),
+                                              partitionings.JoinIndex(),
                                               partitionings.Index(),
                                               partitionings.Singleton()])):
           if not expr.requires_partition_by().is_subpartitioning_of(
diff --git a/sdks/python/apache_beam/dataframe/transforms.py b/sdks/python/apache_beam/dataframe/transforms.py
@@ -375,7 +375,7 @@ def expr_to_stages(expr):
                   expr.requires_partition_by() == stage.partitioning)
         if inputs_by_stage:
           # Take the stage with the largest count.
-          stage = sorted(inputs_by_stage.items(), key=lambda kv: kv[1])[-1][0]
+          stage = max(inputs_by_stage.items(), key=lambda kv: kv[1])[0]
         else:
           stage = None
       else:
