[SPARK-29240][PYTHON] Pass Py4J column instance to support PySpark column in element_at function

[HyukjinKwon]

What changes were proposed in this pull request?

This PR makes element_at in PySpark able to take PySpark Column instances.

Why are the changes needed?

To match with Scala side. Seems it was intended but not working correctly as a bug.

Does this PR introduce any user-facing change?

Yes. See below:

from pyspark.sql import functions as F
x = spark.createDataFrame([([1,2,3],1),([4,5,6],2),([7,8,9],3)],['list','num'])
x.withColumn('aa',F.element_at('list',x.num.cast('int'))).show()

Before:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/.../spark/python/pyspark/sql/functions.py", line 2059, in element_at
    return Column(sc._jvm.functions.element_at(_to_java_column(col), extraction))
  File "/.../spark/python/lib/py4j-0.10.8.1-src.zip/py4j/java_gateway.py", line 1277, in __call__
  File "/.../spark/python/lib/py4j-0.10.8.1-src.zip/py4j/java_gateway.py", line 1241, in _build_args
  File "/.../spark/python/lib/py4j-0.10.8.1-src.zip/py4j/java_gateway.py", line 1228, in _get_args
  File "/.../forked/spark/python/lib/py4j-0.10.8.1-src.zip/py4j/java_collections.py", line 500, in convert
  File "/.../spark/python/pyspark/sql/column.py", line 344, in __iter__
    raise TypeError("Column is not iterable")
TypeError: Column is not iterable

After:

+---------+---+---+
|     list|num| aa|
+---------+---+---+
|[1, 2, 3]|  1|  1|
|[4, 5, 6]|  2|  5|
|[7, 8, 9]|  3|  9|
+---------+---+---+

How was this patch tested?

Manually tested against literal, Python native types, and PySpark column.

[HyukjinKwon]

Otherwise:

flake8 checks failed:
./python/pyspark/sql/functions.py:2059:70: F821 undefined name 'lit'
    return Column(sc._jvm.functions.element_at(_to_java_column(col), lit(extraction)._jc))
                                                                     ^
1     F821 undefined name 'lit'
1
1

[SparkQA]

Test build #111482 has finished for PR 25950 at commit a86d188.

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

[dongjoon-hyun]

+1, LGTM. Thank you, @HyukjinKwon . Merged to master/2.4.