[SPARK-3679] [PySpark] pickle the exact globals of functions

[davies]

function.func_code.co_names has all the names used in the function, including name of attributes. It will pickle some unnecessary globals if there is a global having the same name with attribute (in co_names).

There is a regression introduced by #2144, revert part of changes in that PR.

cc @JoshRosen

[SparkQA]

QA tests have started for PR 2522 at commit dfbccf5.

• This patch merges cleanly.

[JoshRosen]

Based on the JIRA, it looks like this PR description should link to #2144 instead?

[JoshRosen]

Note for other reviewers: this commit mostly reverts #2144's changes to extract_code_globals, so although the changed code is complicated, we were using it before in cloudpickle and should be able to trust that it works.

[davies]

@JoshRosen fixed the description

[JoshRosen]

Based on a read through this PR, it looks like this line is the first place where this function diverges from the pre-#2144 version of cloudpickle.

It looks like the original version of cloudpickle called this code from outside of extract_code_globals, so I guess the old code would only perform one level of recursion when trying to extract globals?

Do you think that adding actual, unbounded recursion could cause problems here? If the "nested function" implies that this only applies to functions defined within other functions, then there aren't cycles in the nesting and therefore shouldn't be cycles that lead to infinite recursion.

[davies]

The original code only perform two levels of functions, the new version can handle multiple levels.

Each level is on function code, which is created by def or lambda, so I think the code cannot be recursive.

[JoshRosen]

This looks good to me, pending Jenkins.

[SparkQA]

QA tests have finished for PR 2522 at commit dfbccf5.

• This patch passes unit tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins//job/SparkPullRequestBuilder/20759/

[JoshRosen]

LGTM; merging this now. Thanks!