Some functions take only Column type as argument about pyspark-stubs

[zero323]

@FurcyPin I am going to close it with #163 - you encounter any similar cases, please let me know.

[zero323]

That's a valid point, thank you for reporting.

Related:

• [DISCUSS][SQL][PySpark] Column name support for SQL functions on Apache Spark Developers List.

[zero323]

For simple unary and binary ((Column, Column) => Column) functions we can identify possible mismatches like this:

from operator import attrgetter, itemgetter
import re

from py4j.protocol import Py4JError
import requests
from toolz.curried import (
    filter, groupby, itemfilter, map, mapcat, pipe, sorted, valfilter, valmap
)
from toolz.functoolz import compose

from pyspark.sql import functions
from pyspark.sql.utils import ParseException

url = (
    "https://raw.githubusercontent.com/apache/spark/"
    "v2.4.3/sql/core/src/main/scala/org/apache/spark/sql/functions.scala"
)
func_pattern = re.compile("^\s*def ([\w_]+)\(([, :\w]*?)\): Column")
argtype_pattern = re.compile("[\w_]+: (\w+)")

def is_candidate(types):
    pref = pipe(types, map(itemgetter(slice(0, 2))), set)
    return (
        (("Column", ) in pref and ("String", ) not in pref) or
        (("Column", "Column") in pref and ("String", "String") not in pref)
    )


def takes_string(item):
    name, types = item
    try:
        f = getattr(functions, name)
        if ("Column", "Column") in types:
            f("foo", "bar")
        else:
            f("foo")
        # Successfully applied string
        return False
    # Incorrect arity
    except TypeError:
        return False
    # No such function in PySpark
    except AttributeError:
        return False
    except ParseException:
        return False
    # Doesn't take String as the first argument
    except Py4JError:
        return True

pipe(
    requests.get(url).text.splitlines(),
    mapcat(func_pattern.findall),
    sorted,
    groupby(itemgetter(0)),
    valmap(compose(
        set,
        map(compose(tuple, argtype_pattern.findall)),
        map(itemgetter(1))
    )),
    valfilter(is_candidate),
    itemfilter(takes_string),
    list
)

which gives in 2.4.3:

['abs',
 'array_repeat',
 'ascii',
 'base64',
 'bitwiseNOT',
 'lower',
 'ltrim',
 'rtrim',
 'trim',
 'unbase64',
 'upper']

but excluding array_repeat (which is a different beast), these are already fixed in 3.0.0.pre0 (master).