Add integer in repeating N rows based on a condition

user_13912414 0
Viewed 71

I am provided with the following DataFrame with panel data containing Dates, Stocks and some feature for each stock.

+----------+------+-------------------+--------------------+
|       _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|
+----------+------+-------------------+--------------------+

I would like to create a new PySpark dataframe containing signals defined as follows: if in 2000-12-01 stock price for XEL > 0, then populatate the next N days for stock XEL with 1, if negative with -1 and i want them to be additive (if in the next trading day price is positive I want to add the signal 1 on the next N days). I want to do this for all the columns.

How can I achieve such result with PySpark DataFrames without iterating through the rows?

An example of the output is the following:

+----------+------+-------------------+--------------------+
|       _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|
+----------+------+-------------------+--------------------+
pyspark apache-spark-sql apache-spark dataframe
1 Answers

You can sum over the signum of the preceding rows:

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|
+----------+------+----------+----+
user_141657300
Related