From 998152f194e7e0e1431d7d0bc043594461782c55 Mon Sep 17 00:00:00 2001 From: Takuya UESHIN Date: Tue, 21 Mar 2023 19:41:01 -0700 Subject: [PATCH 1/2] Add array_prepend to Spark Connect Python client. --- python/pyspark/sql/connect/functions.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/pyspark/sql/connect/functions.py b/python/pyspark/sql/connect/functions.py index f01e333c9d765..1c73bcdc84922 100644 --- a/python/pyspark/sql/connect/functions.py +++ b/python/pyspark/sql/connect/functions.py @@ -1248,6 +1248,13 @@ def array_position(col: "ColumnOrName", value: Any) -> Column: array_position.__doc__ = pysparkfuncs.array_position.__doc__ +def array_prepend(col: "ColumnOrName", value: Any) -> Column: + return _invoke_function("array_prepend", _to_col(col), lit(value)) + + +array_prepend.__doc__ = pysparkfuncs.array_prepend.__doc__ + + def array_remove(col: "ColumnOrName", element: Any) -> Column: return _invoke_function("array_remove", _to_col(col), lit(element)) From a5a86a1f74d8770eadc6ee0f85efeb6c784cbe28 Mon Sep 17 00:00:00 2001 From: Takuya UESHIN Date: Tue, 21 Mar 2023 21:33:06 -0700 Subject: [PATCH 2/2] Test. --- .../sql/tests/connect/test_connect_function.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/pyspark/sql/tests/connect/test_connect_function.py b/python/pyspark/sql/tests/connect/test_connect_function.py index a984bba1b6639..9a9ecc1a6fec4 100644 --- a/python/pyspark/sql/tests/connect/test_connect_function.py +++ b/python/pyspark/sql/tests/connect/test_connect_function.py @@ -1088,6 +1088,20 @@ def test_collection_functions(self): sdf.select(SF.array_append(sdf.a, sdf.f)).toPandas(), ) + # test array_prepend + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, "xyz")).toPandas(), + sdf.select(SF.array_prepend(sdf.a, "xyz")).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, CF.lit("ab"))).toPandas(), + sdf.select(SF.array_prepend(sdf.a, SF.lit("ab"))).toPandas(), + ) + self.assert_eq( + cdf.select(CF.array_prepend(cdf.a, cdf.f)).toPandas(), + sdf.select(SF.array_prepend(sdf.a, sdf.f)).toPandas(), + ) + # test array_insert self.assert_eq( cdf.select(CF.array_insert(cdf.a, -5, "ab")).toPandas(),