31855

sdks/python/apache_beam/dataframe/transforms.py

@@ -395,7 +395,8 @@ def expr_to_stages(expr):
 
       if stage is None:
         # No stage available, compute this expression as part of a new stage.
-        stage = Stage(expr.args(), expr.requires_partition_by())
+        stage = Stage([arg for arg in expr.args() if arg in inputs],
+                      expr.requires_partition_by())
         for arg in expr.args():
           # For each argument, declare that it is also available in
           # this new stage.
@@ -422,6 +423,7 @@ def expr_to_stage(expr):
 
     @_memoize
     def stage_to_result(stage):
+      # print({expr._id: expr_to_pcoll(expr) for expr in stage.inputs})
       return {expr._id: expr_to_pcoll(expr)
               for expr in stage.inputs} | ComputeStage(stage)
 

sdks/python/apache_beam/transforms/core.py

@@ -1461,7 +1461,7 @@ def partition_for(self, element, num_partitions, *args, **kwargs):
 
 def _get_function_body_without_inners(func):
   source_lines = inspect.getsourcelines(func)[0]
-  source_lines = dropwhile(lambda x: x.startswith("@"), source_lines)
+  source_lines = dropwhile(lambda x: x.strip().startswith("@"), source_lines)
   first_def_line = next(source_lines).strip()
   if first_def_line.startswith("def "):
     last_def_line_without_comment = first_def_line.split("#")[0] \