数据框值替换

Question

我想把'yyyy-MM'替换成'yyyy-MM'+'-01'，下面是我的代码，但我没有做到。请注意，我是在databricks上工作的：

from pyspark.sql.functions import col, concat, lit, when


# Show the DataFrame
new_df.show(5)

from pyspark.sql.functions import col, concat, lit, when

# Create new columns with replaced values
new_df = clinicaltrial_df.withColumn(
    'Start_new',
    when(
        col('Start').contains('-'),
        col('Start')
    ).otherwise(
        concat(col('Start'), lit('-01'))
    )
).withColumn(
    'Complete_new',
    when(
        col('Completion').contains('-'),
        col('Completion')
    ).otherwise(
        concat(col('Completion'), lit('-01'))
    )
)

# Show the DataFrame
new_df.show(5)

Answer 1

你的代码是为了给数据框中的 Start 和 Completion 列的值后面加上 '-01'，前提是这些值里还没有包含 '-'。不过，你似乎特别想要处理那些格式为 'yyyy-MM' 的字符串，并确保它们变成 'yyyy-MM-01'。为了做到这一点，你需要找出那些完全符合 'yyyy-MM' 格式的字符串。

你可以使用 PySpark 的 regexp_replace 函数，这个函数可以根据一个规则表达式去查找字符串，并把匹配的部分替换成你指定的内容。在你的情况下，你可以查找那些符合年份和月份格式（'yyyy-MM'）并且不以日期结尾的字符串。然后，把 '-01' 加到这些字符串后面，使它们变成完整的日期格式（'yyyy-MM-01'）。

下面是你可以调整的代码：

from pyspark.sql.functions import regexp_replace

# Adjust the DataFrame
new_df = clinicaltrial_df.withColumn(
    'Start_new',
    regexp_replace('Start', r'^(\d{4}-\d{2})$', concat(col('Start'), lit('-01')))
).withColumn(
    'Complete_new',
    regexp_replace('Completion', r'^(\d{4}-\d{2})$', concat(col('Completion'), lit('-01')))
)

# Show the modified DataFrame
new_df.show(5)

上面的方法确保只有格式完全为 'yyyy-MM' 的字符串会被修改，准确地满足了你的需求，并利用了正则表达式强大的模式匹配能力来实现想要的转换。

示例

下面是针对一些示例数据的解决方案：

from pyspark.sql import SparkSession
from pyspark.sql.functions import regexp_replace, col, concat, lit, when
from pyspark.sql.types import StructType, StructField, StringType

# Initialize SparkSession (not necessary if you're running this in Databricks as it's already initialized)
spark = SparkSession.builder.appName("example").getOrCreate()

# Define schema for the DataFrame
schema = StructType([
    StructField("ID", StringType(), True),
    StructField("Start", StringType(), True),
    StructField("Completion", StringType(), True)
])

# Sample data
data = [
    ("001", "2022-01", "2022-12-31"),
    ("002", "2022-05-01", "2023-05"),
    ("003", "2023", "2023-11"),
    ("004", "2022-07", "2022-09-15")
]

# Create DataFrame
clinicaltrial_df = spark.createDataFrame(data, schema=schema)

# Showing original DataFrame
clinicaltrial_df.show()
# +---+----------+----------+
# | ID|     Start|Completion|
# +---+----------+----------+
# |001|   2022-01|2022-12-31|
# |002|2022-05-01|   2023-05|
# |003|      2023|   2023-11|
# |004|   2022-07|2022-09-15|
# +---+----------+----------+

# Apply the solution
new_df = clinicaltrial_df.withColumn(
    'Start_new',
    regexp_replace('Start', r'^(\d{4}-\d{2})$', concat(col('Start'), lit('-01')))
).withColumn(
    'Complete_new',
    regexp_replace('Completion', r'^(\d{4}-\d{2})$', concat(col('Completion'), lit('-01')))
)

# Show the modified DataFrame
new_df.show(truncate=False)
# +---+----------+----------+----------+------------+
# |ID |Start     |Completion|Start_new |Complete_new|
# +---+----------+----------+----------+------------+
# |001|2022-01   |2022-12-31|2022-01-01|2022-12-31  |
# |002|2022-05-01|2023-05   |2022-05-01|2023-05-01  |
# |003|2023      |2023-11   |2023      |2023-11-01  |
# |004|2022-07   |2022-09-15|2022-07-01|2022-09-15  |
# +---+----------+----------+----------+------------+

注意： 只有在字符串中同时包含年份和月份，并且顺序正确的情况下，日期才会被加上。如果你还有只包含年份的日期（'YYYY'），或者年份和月份的顺序不是 'YYYY-MM'，或者日期没有使用 '-' 作为年份和月份的分隔符（例如 'YYYY/MM'），这些值将不会被修改。

正则表达式模式解析

如果你对正则表达式不太熟悉，下面是对上面代码中使用的模式的解析：

• regexp_replace 使用了一个正则表达式模式 r'^(\d{4}-\d{2})$'：
  • ^ 表示字符串的开始。
  • (\d{4}-\d{2}) 匹配并捕获一个由四个数字（表示年份）和一个连字符后跟两个数字（表示月份）组成的组。
  • $ 表示字符串的结束。