如何在PySpark中基于startswith函数转换DataFrame列并封装函数？

Hey there! Since you're new to PySpark and Python, I'll break this down into two easy-to-follow approaches—one using PySpark's optimized built-in functions (the recommended, faster way) and another using a custom UDF (User Defined Function) as you requested.

First, Let's Set Up a Sample DataFrame #

Let's start with a sample DataFrame matching your example so we can test our solutions:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, regexp_replace

# Initialize Spark session
spark = SparkSession.builder.appName("PrefixReplacement").getOrCreate()

# Sample data matching your example
data = [("abc-text1",), ("def_text2-",)]
df = spark.createDataFrame(data, ["col"])

df.show()
# Output:
# +----------+
# |       col|
# +----------+
# |abc-text1 |
# |def_text2-|
# +----------+

Approach 1: Using PySpark Built-in Functions (Recommended) #

PySpark's built-in functions run natively on the JVM, making them way faster than Python UDFs. Here's how to handle your prefix replacement with when and regexp_replace:

modified_df = df.withColumn(
    "modified_col",
    when(col("col").startswith("abc-"), regexp_replace(col("col"), "^abc-", "abc"))
    .when(col("col").startswith("def_"), regexp_replace(col("col"), "^def_", "def"))
    .otherwise(col("col"))  # Keep original value if no prefix match
)

modified_df.show()
# Output:
# +----------+-----------+
# |       col|modified_col|
# +----------+-----------+
# |abc-text1 |      abc   |
# |def_text2-|      def   |
# +----------+-----------+

For an even simpler approach (since we just need to keep the first 3 characters of matching strings), use substring:

modified_df = df.withColumn(
    "modified_col",
    when(col("col").startswith("abc-"), col("col").substr(1, 3))
    .when(col("col").startswith("def_"), col("col").substr(1, 3))
    .otherwise(col("col"))
)

Approach 2: Custom Python UDF #

If you specifically want a custom function, here's how to create a UDF. Note that UDFs are slower than built-in functions, but they're handy for more complex logic that can't be handled with PySpark's native tools.

First, define your Python function:

def replace_prefix(s):
    if s is None:  # Handle null values gracefully
        return s
    if s.startswith("abc-"):
        return "abc"
    elif s.startswith("def_"):
        return "def"
    else:
        return s  # Return original if no match

Then register it as a PySpark UDF and apply it to your column:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# Register the UDF with a return type of String
prefix_replace_udf = udf(replace_prefix, StringType())

# Apply the UDF to create a new modified column
modified_df = df.withColumn("modified_col", prefix_replace_udf(col("col")))

modified_df.show()
# Same output as the built-in function approach

Quick Tips for Beginners #

• Stick to built-in functions whenever possible—they’re optimized for performance and avoid the overhead of switching between Python and the JVM.
• The otherwise clause ensures any values that don’t match your prefix rules stay unchanged.
• Always add null checks if your column might have missing values, like we did in the UDF example.

内容的提问来源于stack exchange，提问作者Pankaj Kaundal