Spark/Scala如何编写复杂查询遍历DataFrame并新增列

Got it, let's walk through how to solve this problem step by step. You've got a DataFrame with date ranges, and you want to add a new column that expands each range into an array of individual dates using your custom days_in_range function.

First, let's restate your source data clearly:

Original DataFrame (Dataframe 1) #

Since you're working with DataFrames, I assume you're using Spark (the most common framework for this kind of operation). Below are solutions for both PySpark and Scala, matching the function signature you provided.

Solution 1: PySpark Implementation #

Step 1: Define & Register Your Custom Function as a UDF #

Spark can't use plain Python functions directly in DataFrame operations—you need to wrap it into a User-Defined Function (UDF) with the correct return type (an array of strings, since your dates are in yyyyMMdd string format).

First, let's flesh out the days_in_range function (you can replace this with your existing implementation if it's already working):

from datetime import datetime, timedelta
from pyspark.sql import SparkSession
from pyspark.sql.types import ArrayType, StringType

# Initialize Spark session (if not already done)
spark = SparkSession.builder.appName("DateRangeExpansion").getOrCreate()

def days_in_range(start_day: str, end_day: str) -> list[str]:
    # Convert string dates to datetime objects
    start_date = datetime.strptime(start_day, "%Y%m%d")
    end_date = datetime.strptime(end_day, "%Y%m%d")
    
    # Generate all dates in the range
    date_list = []
    current_date = start_date
    while current_date <= end_date:
        date_list.append(current_date.strftime("%Y%m%d"))
        current_date += timedelta(days=1)
    return date_list

# Register the function as a UDF
days_in_range_udf = spark.udf.register("days_in_range_udf", days_in_range, ArrayType(StringType()))

Step 2: Add the New Column to Your DataFrame #

Use withColumn to apply the UDF to each row and create your new date array column:

# Load your existing DataFrame (or use the sample data below to test)
sample_data = [("HIVER 19", "20190602", "20191130"), ("ETE 19", "20181202", "20190601")]
df = spark.createDataFrame(sample_data, ["COLLECTION", "START_DAY_ID", "END_DAY_ID"])

# Add the new date range array column
result_df = df.withColumn("DATE_RANGE_ARRAY", days_in_range_udf(df.START_DAY_ID, df.END_DAY_ID))

# View the result (use truncate=False to see full arrays)
result_df.show(truncate=False)

Solution 2: Scala Implementation #

Matching your provided function signature (def days_in_range(start_day: String...), here's the Scala version:

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types.{ArrayType, StringType}
import java.time.{LocalDate, Period}
import java.time.format.DateTimeFormatter

// Initialize Spark session
val spark = SparkSession.builder.appName("DateRangeExpansion").getOrCreate()

// Define the custom function
def days_in_range(startDay: String, endDay: String): Array[String] = {
  val formatter = DateTimeFormatter.ofPattern("yyyyMMdd")
  val startDate = LocalDate.parse(startDay, formatter)
  val endDate = LocalDate.parse(endDay, formatter)
  val daysBetween = Period.between(startDate, endDate).getDays + 1 // Include start and end dates
  (0 until daysBetween).map(d => startDate.plusDays(d).format(formatter)).toArray
}

// Register as UDF
val daysInRangeUdf = udf(days_in_range _, ArrayType(StringType))

// Load sample DataFrame
val sampleData = Seq(("HIVER 19", "20190602", "20191130"), ("ETE 19", "20181202", "20190601"))
val df = spark.createDataFrame(sampleData).toDF("COLLECTION", "START_DAY_ID", "END_DAY_ID")

// Add the new column
val resultDf = df.withColumn("DATE_RANGE_ARRAY", daysInRangeUdf($"START_DAY_ID", $"END_DAY_ID"))

// Show results
resultDf.show(truncate = false)

Key Notes #

• Date Format Consistency: Make sure the date format used in days_in_range matches the format in your DataFrame (we used yyyyMMdd here, which aligns with your sample data).
• Performance Considerations: If your date ranges are extremely large, generating huge arrays might eat up memory. If you need individual rows instead of an array, add .selectExpr("*", "explode(DATE_RANGE_ARRAY) as INDIVIDUAL_DATE") to the result.

内容的提问来源于stack exchange，提问作者Hela Chikhaoui