Pyspark 中两个数据帧（具有值的年份）之间的插值

Question

如何在表示不同年份（例如 2020 年和 2030 年）数据的两个 PySpark DataFrame 之间实现线性插值，以生成像 2025 年这样的中间年份的新 PySpark DataFrame？两个 DataFrame 具有相同的数值结构。年份具有相同的粒度。

我最初的方法涉及https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.interpolate.html

这是推荐的方式吗？

我不久前写过这个 Pandas 方法，但我需要迁移到 Pyspark，但我很难在 Pandas 中实现相同的方法。

def interpolate_between_years(first: DataFrame, second: DataFrame) -> DataFrame:

    years = [first.index.year[0], second.index.year[0]]
    interpolated_df = (
        pd.concat(
            [first.reset_index(drop=True), second.reset_index(drop=True)],
            keys=years,
            axis=1,
        )
        .T.reindex(np.arange(years[0], years[1] + 1))
        .interpolate()
    )

    return interpolated_df

Answer 1

您可以使用 Spark 上的 Pandas API 进行线性插值：

from pyspark.pandas import DataFrame as psp_DataFrame
import pyspark

def interpolate_between_years(first: psp_DataFrame, second: psp_DataFrame) -> psp_DataFrame:
    year_first = first["year"].iloc[0]
    year_second = second["year"].iloc[0]
    
    years_range = range(year_first, year_second + 1)
    interpolated_df = psp_DataFrame({"year": years_range})

    # Interpolate values for each column
    for col in first.columns:
        if col != "year":
            interpolated_df[col] = first[col].iloc[0] + (second[col].iloc[0] - first[col].iloc[0]) / (year_second - year_first) * (interpolated_df["year"] - year_first)

    return interpolated_df


spark = pyspark.sql.SparkSession.builder.getOrCreate()

data_2020 = [(2020, 10, 20), (2021, 11, 21), (2022, 12, 22)]
df_2020 = psp_DataFrame(data_2020, columns=["year", "value1", "value2"])

data_2030 = [(2030, 40, 50), (2031, 41, 51), (2032, 42, 52)]
df_2030 = psp_DataFrame(data_2030, columns=["year", "value1", "value2"])

interpolated_df = interpolate_between_years(df_2020, df_2030)
pandas_df = interpolated_df.to_pandas()
print(pandas_df)

输出：

    year  value1  value2
0   2020    10.0    20.0
1   2021    13.0    23.0
2   2022    16.0    26.0
3   2023    19.0    29.0
4   2024    22.0    32.0
5   2025    25.0    35.0
6   2026    28.0    38.0
7   2027    31.0    41.0
8   2028    34.0    44.0
9   2029    37.0    47.0
10  2030    40.0    50.0

从此

interpolated_df

，您可以获取任何中间年份的数据，例如 2025 年、2026 年等。