Skip to content

[SPARK-48766][PYTHON] Document the behavior difference of extraction between element_at and try_element_at #47161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zhengruifeng wants to merge 2 commits into from
Closed

[SPARK-48766][PYTHON] Document the behavior difference of extraction between element_at and try_element_at #47161

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
970b37c
init
zhengruifeng Jul 1, 2024
8365824
nit
zhengruifeng Jul 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions python/pyspark/sql/functions/builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14098,10 +14098,13 @@ def element_at(col: "ColumnOrName", extraction: Any) -> Column:
Notes
-----
The position is not zero based, but 1 based index.
If extraction is a string, :meth:`element_at` treats it as a literal string,
while :meth:`try_element_at` treats it as a column name.
Comment on lines +14101 to +14102
Copy link
Contributor

@allisonwang-db allisonwang-db Jul 1, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this behavior difference intentional? Is it consistent with the SQL functions element_at and try_element_at?

Copy link
Contributor Author

@zhengruifeng zhengruifeng Jul 1, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the SQL side also treat it as a literal:

In [2]: spark.sql("SELECT ELEMENT_AT(MAP('a', 'b'), 'a')").show()
+------------------------+
|element_at(map(a, b), a)|
+------------------------+
|                       b|
+------------------------+


See Also
--------
:meth:`get`
:meth:`try_element_at`

Examples
--------
Expand Down Expand Up @@ -14148,6 +14151,17 @@ def element_at(col: "ColumnOrName", extraction: Any) -> Column:
+-------------------+
| NULL|
+-------------------+

Example 5: Getting a value from a map using a literal string as the key

>>> from pyspark.sql import functions as sf
>>> df = spark.createDataFrame([({"a": 1.0, "b": 2.0}, "a")], ['data', 'b'])
>>> df.select(sf.element_at(df.data, 'b')).show()
+-------------------+
|element_at(data, b)|
+-------------------+
| 2.0|
+-------------------+
"""
return _invoke_function_over_columns("element_at", col, lit(extraction))

Expand All @@ -14172,6 +14186,17 @@ def try_element_at(col: "ColumnOrName", extraction: "ColumnOrName") -> Column:
extraction :
index to check for in array or key to check for in map

Notes
-----
The position is not zero based, but 1 based index.
If extraction is a string, :meth:`try_element_at` treats it as a column name,
while :meth:`element_at` treats it as a literal string.

See Also
--------
:meth:`get`
:meth:`element_at`

Examples
--------
Example 1: Getting the first element of an array
Expand Down Expand Up @@ -14228,6 +14253,17 @@ def try_element_at(col: "ColumnOrName", extraction: "ColumnOrName") -> Column:
+-----------------------+
| NULL|
+-----------------------+

Example 6: Getting a value from a map using a column name as the key

>>> from pyspark.sql import functions as sf
>>> df = spark.createDataFrame([({"a": 1.0, "b": 2.0}, "a")], ['data', 'b'])
>>> df.select(sf.try_element_at(df.data, 'b')).show()
Copy link
Contributor

@panbingkun panbingkun Jul 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to add another example below this?

df.select(sf.try_element_at(df.data, df.b)).show()

Because of the grammar scene above, I took a long time to understand its intended meaning, which is so obscure, 😄

+-----------------------+
|try_element_at(data, b)|
+-----------------------+
| 1.0|
+-----------------------+
"""
return _invoke_function_over_columns("try_element_at", col, extraction)

Expand Down