根据条件在重复 N 行中添加整数

Question

我收到了以下 DataFrame，其中包含包含日期、股票和每只股票的一些特征的面板数据。

+----------+------+-------------------+--------------------+
|       _c0|symbol|         stockPrice|                roic|
+----------+------+-------------------+--------------------+
|2000-12-01|   MTB|                0.0|                 0.0|
|2000-12-01|   XEL| 0.5599999999999987| 0.03831091180429823|
|2001-03-01|  ADBE|-1.8799999999999968| 0.16664239102772385|
|2001-03-01|   ALL|  3.020000000000003|0.015843499622776227|
|2001-03-01|   XEL|  5.199999999999998| 0.01718300080156956|
|2001-03-01|   BMY| -6.020000000000002|  0.1386651323360184|
|2001-03-01|   CAT| 2.9000000000000017|-0.01871476386801...|
|2001-03-01|   MRK| -5.030000000000001| 0.09202762084118016|
|2001-03-01|   MTB|              21.55|  0.4299814503261493|
|2001-03-01|   PVH| 2.1500000000000004|-0.08486335480927014|
|2001-03-01|   BBY|               4.75| 0.27611001898177634|
|2001-06-01|  ADBE| 17.610000000000007|-0.02235471561675301|
|2001-06-01|   BMY|                0.0|  0.1386651323360184|
|2001-06-01|   ALK|               7.65| 0.03139856533358338|
|2001-06-01|   CAT| 1.6799999999999995|0.002447684500232757|
|2001-06-01|   HUM|               0.25|-0.05092018990508895|
|2001-06-01|   ALL| -7.170000000000003|-0.01125850244811...|
|2001-06-01|  EXPD|               2.74| 0.09279634515913428|
|2001-06-01|  IDXX|              -0.92| 0.09939743655211153|
|2001-06-01|   BBY| 13.420000000000002|-0.21821519283350074|
+----------+------+-------------------+--------------------+

我想创建一个新的 PySpark 数据框，其中包含如下定义的信号：如果在 2000 年 12 月 1 日，XEL 的股票价格 > 0，则在接下来的 N 天为股票 XEL 填充 1，如果为负，则使用 -1 和我希望它们是附加的（如果在下一个交易日价格是正的，我想在接下来的 N 天添加信号 1）。我想对所有列都这样做。

如何在不遍历行的情况下使用 PySpark DataFrames 实现这样的结果？

输出示例如下：

+----------+------+-------------------+--------------------+
|       _c0|symbol|         stockPrice|                roic|
+----------+------+-------------------+--------------------+
|2000-12-01|   MTB|                0.0|                 0.0|
|2000-12-01|   XEL|                  1|                   1|
|2001-03-01|  ADBE|                 -1|                   1|
|2001-03-01|   ALL|                  1|                   1|
|2001-03-01|   XEL|                  2|                   2|
|2001-03-01|   BMY|                 -1|                   1|
|2001-03-01|   CAT|                  1|                  -1|
|2001-03-01|   MRK|                 -1|                   1|
|2001-03-01|   MTB|                  1|                   1|
|2001-03-01|   PVH|                  1|                  -1|
|2001-03-01|   BBY|                  1|                   1|
+----------+------+-------------------+--------------------+

Answer 1

你可以对前面几行的符号求和：

from pyspark.sql import functions as F, Window

w = Window.partitionBy('symbol').orderBy('_c0')

result = df.select(
    '_c0',
    'symbol', 
    F.sum(F.signum('stockPrice')).over(w).alias('stockPrice'), 
    F.sum(F.signum('roic')).over(w).alias('roic')
).orderBy('_c0', 'symbol')

result.show()
+----------+------+----------+----+
|       _c0|symbol|stockPrice|roic|
+----------+------+----------+----+
|2000-12-01|   MTB|       0.0| 0.0|
|2000-12-01|   XEL|       1.0| 1.0|
|2001-03-01|  ADBE|      -1.0| 1.0|
|2001-03-01|   ALL|       1.0| 1.0|
|2001-03-01|   BBY|       1.0| 1.0|
|2001-03-01|   BMY|      -1.0| 1.0|
|2001-03-01|   CAT|       1.0|-1.0|
|2001-03-01|   MRK|      -1.0| 1.0|
|2001-03-01|   MTB|       1.0| 1.0|
|2001-03-01|   PVH|       1.0|-1.0|
|2001-03-01|   XEL|       2.0| 2.0|
|2001-06-01|  ADBE|       0.0| 0.0|
|2001-06-01|   ALK|       1.0| 1.0|
|2001-06-01|   ALL|       0.0| 0.0|
|2001-06-01|   BBY|       2.0| 0.0|
|2001-06-01|   BMY|      -1.0| 2.0|
|2001-06-01|   CAT|       2.0| 0.0|
|2001-06-01|  EXPD|       1.0| 1.0|
|2001-06-01|   HUM|       1.0|-1.0|
|2001-06-01|  IDXX|      -1.0| 1.0|
+----------+------+----------+----+