Recursion error when using UDFs in pyspark

[skasai5296]

Environment data

• Language Server version: v2020.12.1
• OS and version: macOS Catalina 10.15.7
• Python version (& distribution if applicable, e.g. Anaconda): 3.6.9, in pipenv

Expected behaviour

Pylance works

Actual behaviour

Script runs without errors, but Pylance doesn't work when using the UDF.
Commenting out the get_ids function returns reportUndefinedVariable on the commented out function without the RecursionError, which is expected.

Code Snippet

import pyspark
from pyspark import sql
from pyspark.sql import functions as F, types as T


@F.udf(returnType=T.ArrayType(T.StringType()))
def get_ids():
    return ["a", "b", "c", "d"]


def main():
    spark = pyspark.SparkContext(appName="sample")
    session = sql.SparkSession(spark)
    df = session.createDataFrame([("hoge", 10), ("fuga", 100)])
    df = df.select(get_ids().alias("ids"))
    df.show(5)


if __name__ == "__main__":
    main()

An internal error occurred while type checking file \"/Users/xxx/sample.py\": RangeError: Maximum call stack size exceeded
    at canAssignType (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15557:14)
    at assignTypeToTypeVar (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15478:21)
    at canAssignType (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15634:24)
    at verifyTypeArgumentsAssignable (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15286:30)
    at canAssignClassWithTypeArgs (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15160:18)
    at canAssignClass (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:14998:24)
    at canAssignType (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15961:22)
    at assignTypeToTypeVar (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15521:18)
    at canAssignType (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:15634:24)
    at zn (/Users/xxx/.vscode/extensions/ms-python.vscode-pylance-2020.12.1/dist/pyright-internal/src/analyzer/typeEvaluator.ts:17268:25)

[erictraut]

I'm not able to repro the problem with the above sample with just pyspark installed. I was able to repro it with pyspark-stubs installed. I presume you've installed the stubs.

The cause of the overflow is due to what is arguably a bug in the pyspark stubs. The Column class is annotated to include a __getattr__ method that returns a Column object. When you use a Column object as the left side of a call expression (as in get_ids(), Pylance looks for a __call__ attribute on the object. There is no __call__ attribute, but there is a __getattr__. This returns a Column object, which is then called. This results in infinite recursion.

I've added protection against this case by ignoring any __call__ attribute (or its proxy returned by __getattr__) that is not a function. That will avoid the infinite recursion in the type analyzer. This fix will be in the next release.

[skasai5296]

Thank you @erictraut for your rocket-fast response, I'll look forward to the next release.

[jakebailey]

This issue has been fixed in version 2020.12.2, which we've just released. You can find the changelog here: https://github.com/microsoft/pylance-release/blob/master/CHANGELOG.md#2020122-11-december-2020