Pyspark: Filter dataframe and apply function to offset time

user_6058499 0
Viewed 376

I have a dataframe like this:

import time
import datetime
import pandas as pd

df = pd.DataFrame({'Number': ['1', '2', '1', '1'],
                   'Letter': ['A', 'A', 'B', 'A'],
                   'Time': ['2019-04-30 18:15:00', '2019-04-30 18:15:00', '2019-04-30 18:15:00', '2019-04-30 18:15:00'],
                   'Value': [30, 30, 30, 60]})

df['Time'] = pd.to_datetime(df['Time'])


  Number Letter                Time  Value
0      1      A 2019-04-30 18:15:00     30
1      2      A 2019-04-30 18:15:00     30
2      1      B 2019-04-30 18:15:00     30
3      1      A 2019-04-30 18:15:00     60

I would like to do something similar in Pyspark as I do in Pandas where I filter on a specific set of data:

#: Want to target only rows where the Number = '1' and the Letter is 'A'.

target_df = df[
    (df['Number'] == '1') &
    (df['Letter'] == 'A')
]

And apply a change to a value based on another column:

#: Loop over these rows and subtract the offset value from the Time.
for index, row in target_df.iterrows():
    offset = row['Value']
    df.loc[index, 'Time'] = row['Time'] - datetime.timedelta(seconds=row['Value'])

To get a final output like so:

  Number Letter                Time  Value
0      1      A 2019-04-30 18:14:30     30
1      2      A 2019-04-30 18:15:00     30
2      1      B 2019-04-30 18:15:00     30
3      1      A 2019-04-30 18:14:00     60

What is the best way to go about this in Pyspark? I was thinking something along the lines of this:

pyspark_df = spark.createDataFrame(df)

pyspark_df.withColumn('new_time', F.when(
    F.col('Number') == '1' & F.col('Letter' == 'A'), F.col('Time') - datetime.timedelta(seconds=(F.col('Value')))).otherwise(
    F.col('Time')))

But that doesn't seem to work for me.

pyspark-dataframes pyspark apache-spark-sql
1 Answers

You can try with unix timestamp:

import pyspark.sql.functions as F


cond_val = (F.when((F.col("Number")==1)&(F.col("Letter")=="A")
      ,F.from_unixtime(F.unix_timestamp(F.col("Time"))-F.col("Value")))
      .otherwise(F.col("Time")))

df.withColumn("Time",cond_val).show()

+------+------+-------------------+-----+
|Number|Letter|               Time|Value|
+------+------+-------------------+-----+
|     1|     A|2019-04-30 18:14:30|   30|
|     2|     A|2019-04-30 18:15:00|   30|
|     1|     B|2019-04-30 18:15:00|   30|
|     1|     A|2019-04-30 18:14:00|   60|
+------+------+-------------------+-----+

Just an addition, you dont need iterrows in pandas, just do:

c = df['Number'].eq(1) & df['Letter'].eq('A')
df.loc[c,'Time'] = df['Time'].sub(pd.to_timedelta(df['Value'],unit='s'))
#or faster
#df['Time'] = np.where(c,df['Time'].sub(pd.to_timedelta(df['Value'],unit='s'))
                                                                  #,df['Time'])
user_98406370
Related