Creating a function in PySpark to return column

Hernandes Matias Junior 0
Viewed 30

I have a function to which I need to provide a dataframe and a column. It extracts a string using regex pattern and return is a new dataframe.

However, I want the return to be just the result of regex, just the column.

def extract_strings(dataframe_selected, column_selected):
    dataframe2 = dataframe_selected.withColumn("strings", F.regexp_extract(dataframe_selected[f"{column_selected}"], r"([a-zA-Z]+)", 0))
    return dataframe2

I would like to create a function which I could use like

df['new_column'] = extract_strings(df, 'text')
function pyspark apache-spark regex python
1 Answers

You can create a function which takes a column and returns a column. But you will need to have a dataframe to put that column in, as columns don't "live" on their own, they need a dataframe to contain them.

Function:

from pyspark.sql import functions as F

def extract_strings(column_selected):
    return F.regexp_extract(column_selected, r"([a-zA-Z]+)", 0)

Test:

df = spark.createDataFrame([("abc1",)], ["col_name"])

df.withColumn("new_column", extract_strings("col_name")).show()
# +--------+----------+
# |col_name|new_column|
# +--------+----------+
# |    abc1|       abc|
# +--------+----------+
ZygD0
Related