PySpark Window not working with specified integer bounds

Question

I'm trying to determine whether particular processing jobs experienced a negative outcome in their next five runs. Here's the setup of my data:

from pyspark.sql import SQLContext, functions as func
from pyspark.sql.window import Window
import datetime

job_history_df = sqlContext.createDataFrame(
    [
        ('A', 'X', datetime.datetime.strptime('2018-01-02 19:00:13.0', '%Y-%m-%d %H:%M:%S.0')),
        ('A', 'X', datetime.datetime.strptime('2018-01-03 19:00:09.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'X', datetime.datetime.strptime('2018-01-04 19:00:24.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'X', datetime.datetime.strptime('2018-01-05 19:00:21.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'X', datetime.datetime.strptime('2018-01-06 19:00:33.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'Y', datetime.datetime.strptime('2018-01-08 19:00:12.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'Y', datetime.datetime.strptime('2018-01-09 19:00:22.0', '%Y-%m-%d %H:%M:%S.0')),
        ('A', 'Y', datetime.datetime.strptime('2018-01-10 19:00:21.0', '%Y-%m-%d %H:%M:%S.0')),
        ('S', 'Y', datetime.datetime.strptime('2018-01-10 19:00:23.0', '%Y-%m-%d %H:%M:%S.0')),

    ],
    ['jhr_status', 'ajb_name', 'jhr_run_date']
)

def dt_to_timestamp():
    def _dt_to_timestamp(dt):
        return int(dt.timestamp() * 1000)
    return func.udf(_dt_to_timestamp)

job_history_df = job_history_df.withColumn('jhr_run_date_ts', dt_to_timestamp()(func.col('jhr_run_date')).cast('long'))
job_history_df = job_history_df.withColumn('was_abend', func.when(job_history_df['jhr_status'] == 'A', 1).otherwise(0))

Here's what job_history_df looks like:

>>> job_history_df.show(20, False)
+----------+--------+---------------------+---------------+---------+
|jhr_status|ajb_name|jhr_run_date         |jhr_run_date_ts|was_abend|
+----------+--------+---------------------+---------------+---------+
|A         |X       |2018-01-02 19:00:13.0|1514941213000  |1        |
|A         |X       |2018-01-03 19:00:09.0|1515027609000  |1        |
|S         |X       |2018-01-04 19:00:24.0|1515114024000  |0        |
|S         |X       |2018-01-05 19:00:21.0|1515200421000  |0        |
|S         |X       |2018-01-06 19:00:33.0|1515286833000  |0        |
|S         |Y       |2018-01-08 19:00:12.0|1515459612000  |0        |
|S         |Y       |2018-01-09 19:00:22.0|1515546022000  |0        |
|A         |Y       |2018-01-10 19:00:21.0|1515632421000  |1        |
|S         |Y       |2018-01-10 19:00:23.0|1515632423000  |0        |
+----------+--------+---------------------+---------------+---------+

>>> job_history_df.dtypes
[('jhr_status', 'string'), ('ajb_name', 'string'), ('jhr_run_date', 'timestamp'), ('jhr_run_date_ts', 'bigint'), ('was_abend', 'int')]

Next, I'll create my Window:

base_job_window = Window().partitionBy('ajb_name').orderBy('jhr_run_date_ts')

Next we'll specify the range I want to sum over:

n_next_runs = 5

next_n_runs_window = base_job_window.rangeBetween(1, n_next_runs)
job_history_df = job_history_df.withColumn('n_abends_next_n_runs', func.sum('was_abend').over(next_n_runs_window))

Let's see what we get:

>>> job_history_df.show(20, False)
+----------+--------+---------------------+---------------+---------+--------------------+
|jhr_status|ajb_name|jhr_run_date         |jhr_run_date_ts|was_abend|n_abends_next_n_runs|
+----------+--------+---------------------+---------------+---------+--------------------+
|S         |Y       |2018-01-08 19:00:12.0|1515459612000  |0        |null                |
|S         |Y       |2018-01-09 19:00:22.0|1515546022000  |0        |null                |
|A         |Y       |2018-01-10 19:00:21.0|1515632421000  |1        |null                |
|S         |Y       |2018-01-10 19:00:23.0|1515632423000  |0        |null                |
|A         |X       |2018-01-02 19:00:13.0|1514941213000  |1        |null                |
|A         |X       |2018-01-03 19:00:09.0|1515027609000  |1        |null                |
|S         |X       |2018-01-04 19:00:24.0|1515114024000  |0        |null                |
|S         |X       |2018-01-05 19:00:21.0|1515200421000  |0        |null                |
|S         |X       |2018-01-06 19:00:33.0|1515286833000  |0        |null                |
+----------+--------+---------------------+---------------+---------+--------------------+

That is strange. The output for n_abends_next_n_runs should be 1s all the way down, I believe. What about if we sum all previous failures?

all_previous_window = base_job_window.rangeBetween(Window.unboundedPreceding, -1)
job_history_df = job_history_df.withColumn('n_abends_to_pt', func.sum('was_abend').over(all_previous_window))

This gives the right result:

+----------+--------+---------------------+---------------+---------+--------------------+--------------+
|jhr_status|ajb_name|jhr_run_date         |jhr_run_date_ts|was_abend|n_abends_next_n_runs|n_abends_to_pt|
+----------+--------+---------------------+---------------+---------+--------------------+--------------+
|S         |Y       |2018-01-08 19:00:12.0|1515459612000  |0        |null                |null          |
|S         |Y       |2018-01-09 19:00:22.0|1515546022000  |0        |null                |0             |
|A         |Y       |2018-01-10 19:00:21.0|1515632421000  |1        |null                |0             |
|S         |Y       |2018-01-10 19:00:23.0|1515632423000  |0        |null                |1             |
|A         |X       |2018-01-02 19:00:13.0|1514941213000  |1        |null                |null          |
|A         |X       |2018-01-03 19:00:09.0|1515027609000  |1        |null                |1             |
|S         |X       |2018-01-04 19:00:24.0|1515114024000  |0        |null                |2             |
|S         |X       |2018-01-05 19:00:21.0|1515200421000  |0        |null                |2             |
|S         |X       |2018-01-06 19:00:33.0|1515286833000  |0        |null                |2             |
+----------+--------+---------------------+---------------+---------+--------------------+--------------+

What could be the issue with specify integer bounds instead of using Window.unboundedPreceding or Window.unboundedFollowing?

For reference, I am running Spark version 2.1.1.2.6.2.14-5 on a RHEL6 VM.

More analysis

Digging into this further, I thought to check out whether "regular-old" SQL would work:

job_history_df.registerTempTable('table')

job_history_df = sqlContext.sql(
    '''
    SELECT
        *,
        SUM(was_abend) OVER (PARTITION BY ajb_name ORDER BY jhr_run_date_ts ROWS BETWEEN 5 PRECEDING AND 1 PRECEDING) AS abends_last_5_runs
    FROM table
    '''
)

In fact, it does!

>>> job_history_df.show(20, False)
+----------+--------+---------------------+---------------+---------+--------------------+------------------+
|jhr_status|ajb_name|jhr_run_date         |jhr_run_date_ts|was_abend|n_abends_next_n_runs|abends_last_5_runs|
+----------+--------+---------------------+---------------+---------+--------------------+------------------+
|S         |Y       |2018-01-08 19:00:12.0|1515459612000  |0        |null                |null              |
|S         |Y       |2018-01-09 19:00:22.0|1515546022000  |0        |null                |0                 |
|A         |Y       |2018-01-10 19:00:21.0|1515632421000  |1        |null                |0                 |
|S         |Y       |2018-01-10 19:00:23.0|1515632423000  |0        |null                |1                 |
|A         |X       |2018-01-02 19:00:13.0|1514941213000  |1        |null                |null              |
|A         |X       |2018-01-03 19:00:09.0|1515027609000  |1        |null                |1                 |
|S         |X       |2018-01-04 19:00:24.0|1515114024000  |0        |null                |2                 |
|S         |X       |2018-01-05 19:00:21.0|1515200421000  |0        |null                |2                 |
|S         |X       |2018-01-06 19:00:33.0|1515286833000  |0        |null                |2                 |
+----------+--------+---------------------+---------------+---------+--------------------+------------------+

While this still doesn't solve the issue of the pure Spark-SQL attempt from attempt, it does make my day at work a lot easier tomorrow :)

Answer 1

Your spark-sql works because you have used rowsBetween in the query instead of rangeBetween. You have used rangeBetween in your first two tries whereas you have used rowsBetween in the query.

The syntax format looks the same for both rowsBetween and rangeBetween, but they perform completely different ways.

Let me illustrate with an example by creating a dataframe and use the same logic as you have used with rowsBetween and rangeBetween but instead of sum, collect_list is used so that it is clear which and how many rows are being considered

lets say you have a dataframe as

df = spark.createDataFrame([('X', -2),
                            ('X', 0),
                            ('X', 2),
                            ('X', 3),
                            ('X', 21),
                            ('X', 1)], ('col1', 'col2'))
+----+----+
|col1|col2|
+----+----+
|   X|  -2|
|   X|   0|
|   X|   2|
|   X|   3|
|   X|  21|
|   X|   1|
+----+----+

rowsBetween

base_job_window = Window().partitionBy('col1').orderBy('col2')
n_next_runs = 4
next_n_runs_window = base_job_window.rowsBetween(2, n_next_runs)

df.withColumn('col2_next_n_runs', func.collect_list('col2').over(next_n_runs_window)).show(truncate=False)

+----+----+----------------+
|col1|col2|col2_next_n_runs|
+----+----+----------------+
|X   |-2  |[1, 2, 3]       |
|X   |0   |[2, 3, 21]      |
|X   |1   |[3, 21]         |
|X   |2   |[21]            |
|X   |3   |[]              |
|X   |21  |[]              |
+----+----+----------------+

As you can see that 2nd 3rd and 4th rows from the current row are collected as the values in rowsBetween are 2 and 4

rangeBetween

base_job_window = Window().partitionBy('col1').orderBy('col2')
n_next_runs = 4
next_n_runs_window = base_job_window.rangeBetween(2, n_next_runs)

df.withColumn('col2_next_n_runs', func.collect_list('col2').over(next_n_runs_window)).show(truncate=False)

+----+----+----------------+
|col1|col2|col2_next_n_runs|
+----+----+----------------+
|X   |-2  |[0, 1, 2]       |
|X   |0   |[2, 3]          |
|X   |1   |[3]             |
|X   |2   |[]              |
|X   |3   |[]              |
|X   |21  |[]              |
+----+----+----------------+

As you can see that the values of rangeBetween are added to the value of the column used in orderBy in the current row and all the values that lie in between that range in col2 are collected.

It would be more clear if you read Introducing Window Functions in Spark SQL

Moreover I am copy pasting part of that article