feat(pyspark): support windowing functions in Pyspark backend

[chloeh13q]

Is your feature request related to a problem?

Windowing functions are not supported in Pyspark backend right now.

Describe the solution you'd like

I'd like to have support for windowing functions in the Pyspark backend: https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html#types-of-time-windows

The traditional SQL windowing with over() is not supported in Spark Structured Streaming (https://stackoverflow.com/questions/55591774/how-to-use-a-non-time-based-window-with-spark-data-streaming-structure)

What version of ibis are you running?

main

What backend(s) are you using, if any?

Pyspark

Code of Conduct

• I agree to follow this project's Code of Conduct

[chloeh13q]

I looked at streaming window aggregations for Flink, Spark, RW, Decodable, Materialize.

Note that Spark also has something called window functions, which corresponds to the over() syntax in Ibis. It's not supported in streaming.

Syntax for streaming windows in Spark

window_agg = session.sql("""
SELECT
    window.start, window.end, provinceId, sum(payAmount) as totalPayAmount
    FROM streaming_df
    GROUP BY provinceId, window(createTime, '1 hour', '30 minutes')
""")

Spark supports three types of time windows: tumbling (fixed), sliding, and session.

We will address session windows separately because 1) it's not currently supported in Ibis' Flink backend and 2) they have different characteristics and use a different API in Spark.

Tumbling and sliding windows both use the window function in Spark. slide_duration is an optional param. The window function takes in both a timestamp column and a time window column. The windowing function produces window.start and window.end. It's also possible to call window_time on top of the window to get a timestamp representation of the time window, e.g.,

SELECT a, window.start as start, window.end as end, window_time(window), cnt FROM (SELECT a, window, count(*) as cnt FROM VALUES ('A1', '2021-01-01 00:00:00'), ('A1', '2021-01-01 00:04:30'), ('A1', '2021-01-01 00:06:00'), ('A2', '2021-01-01 00:01:00') AS tab(a, b) GROUP by a, window(b, '5 minutes') ORDER BY a, window.start);
+---+-------------------+-------------------+--------------------+---+
|  a|              start|                end| window_time(window)|cnt|
+---+-------------------+-------------------+--------------------+---+
| A1|2021-01-01 00:00:00|2021-01-01 00:05:00|2021-01-01 00:04:...|  2|
| A1|2021-01-01 00:05:00|2021-01-01 00:10:00|2021-01-01 00:09:...|  1|
| A2|2021-01-01 00:00:00|2021-01-01 00:05:00|2021-01-01 00:04:...|  1|
+---+-------------------+-------------------+--------------------+---+

Syntax in other streaming engines/databases

• Decodable
• Flink
• RW
• Materialize

Flink SQL and Decodable SQL both use table-valued functions.
RW uses tumble() and hop() functions. The result of a time window function is a table in which each row carries data for a time window. A time window function extends the schema of the original table with two new columns, window_start and window_end, which indicate the start and end of time windows respectively.
Materialize has something called a date_bin function.

Proposal for implementation

Ibis has a window function but I think the semantics are too different. We also introduced a window_by method that returns a WindowedTable object when we implemented windowing TVFs in Flink.

In Flink, windowing TVFs return tables. Spark's window() function returns a single column. The output column will be a struct called 'window' by default with the nested columns 'start' and 'end', where 'start' and 'end' will be of pyspark.sql.types.TimestampType. You can use this column like you would any other column - it doesn't have to be combined with a groupby, e.g.,

streaming_df = streaming_df.withColumn(
    'window', 
    window(streaming_df.createTime, '1 hour', '30 minutes')
)

This outputs records like

{"createTime":"2024-04-10T20:55:07.629-07:00","orderId":1712783621,"payAmount":64630.64838013934,"payPlatform":0,"provinceId":3,"window":{"start":"2024-04-10T20:00:00.000-07:00","end":"2024-04-10T21:00:00.000-07:00"}}

Usage of the window_by method in Ibis:

expr = (
    simple_table.window_by(time_col=simple_table.i)
    .tumble(window_size=ibis.interval(minutes=15))
    .group_by(["window_start", "window_end", "g"])
    .aggregate(mean=_.d.mean())
)

We can have .window_by().tumble() just return a table with all the columns + an additional window column in Spark. Then group by behaves the same way. Basically deconstructing the query into two steps, i.e.,

# Dataframe API
streaming_df = streaming_df.withColumn(
    'window', 
    window(streaming_df.createTime, '1 minute', '30 seconds')
)

streaming_df = streaming_df.groupBy("window").count()

# equivalent SQL
streaming_df = session.sql("""
SELECT window, provinceId, sum(payAmount) FROM
  (SELECT *, window(createTime, '1 minute', '30 seconds') FROM streaming_df) AS t0
GROUP BY t0.window, t0.provinceId
""")

What we will need to do to implement this approach:

• Make the 3 TVF ops more generic windowing ops, i.e., remove the TVF from the name
• How do we handle the schema of the windowing window output? This is backend-specific. Flink adds 3 additional columns to the table (window_start, window_end, window_time); Spark adds a single column (window) with a nested struct ({'start':, 'end':,}); RW adds two columns (window_start and window_end). Maybe we should just deferring the handling of this schema to the user.
• Implement compilation rules for the windowing ops in Spark backend

Note

• You can't have a column named window in your upstream source table
• Spark is less strict about timestamp precision than Flink

[chloeh13q]

We had a discussion around this. There are two options that I have considered:

1. Something along the lines of the TimestampBucket operation.

I think Spark's window() function is sort of similar to the .bucket() method in Ibis.

Example usage of bucket():

mean_by_bucket = (
     t.group_by(t.ts.bucket(minutes=5).name("bucket"))
     .agg(mean=_.val.mean())
     .order_by("bucket")
)

It feels like a streaming equivalent of the TimestampBucket op. But the buckets here are non-overlapping, i.e., it cannot express sliding windows. This is an alternative syntax for windowing that we may want to explore.

2. Use the window_by API as I have mentioned above. We discussed whether this seems to be a reasonable abstraction, and an important part of the consideration there being "capturing user intent in a generic way and translating that to the backends", as opposed to the other way around. While it does make specific operations more expensive (particularly if the user just wants the windows column, due to the constraint here we will basically build an entire table, then select that column, which may be slower depending on how that query is being optimized underneath the hood), I think it's reasonable to assume that these these usage patterns being affected are rare (i.e., it's rare for the user to only want the window column as output, because that's basically generating a floor/ceiling and it doesn't make much sense if it's not connected to the rest of the data).

The conclusion was to potentially implement both in the Spark streaming backend, with window_by being more of just syntactic sugar but a more natural way for people coming from Spark/Flink/etc. But I also want to call out that the benefit of implementing bucket may be limited because it only support 1 out of the 4 window types.

Separately, I'm still debating on whether it makes sense for us to build the "group by" into the windowing, the way that it works in pandas. In Spark/Flink/other streaming engines, all that the window does is that it creates the time buckets. In order to actually do anything meaningful, you need to group by the time buckets. Arguably that seems obvious and perhaps a bit of boilerplate, but there is more flexibility, say, if you only want to group by the start of the window or the end of the window. On the contrary, a pandas-like API that does group by as a part of the windowing does not provide that flexibility.

I'm leaning towards avoiding making large changes to the window_by API right now and continue to collect feedback.

[jcrist]

We can have .window_by().tumble() just return a table with all the columns + an additional window column in Spark. Then group by behaves the same way.

This is what I would do, except I'd make it match the schema of the existing window_by operations. We can generate output tables with the same schema that flink generates by adding an extra projection to create the window_start/window_end/window_time columns. The window_by functionality as exists should map well enough onto how spark does things, and I'd expect/hope that their optimizer should elide creating the extra projected columns if they're not referenced later on, so I wouldn't expect coercing spark's output to match this schema to come with a large perf cost.

Make the 3 TVF ops more generic windowing ops, i.e., remove the TVF from the name

I don't think we necessarily need to rename them - conceptually they're still table-valued functions (take in a table and some options, return a table). Not opposed to renaming though.

Something along the lines of the TimestampBucket operation.

I added TimestampBucket originally to be able to implement window_by for other non-streaming backends (I still intend to finish this up). Sparks window is a superset of this functionality, so while we could implement bucket using spark's window, IMO we should aim towards supporting window_by in spark first.

[chloeh13q]

@jcrist Makes sense, I think we're on the same page. Just wondering: what is the motivation behind matching the schema of the existing window_by operations? I'm not opposed to it and I think it's one of the two options we have, but just wondering if what your considerations there are.

[jcrist]

Mainly I think we want to ensure the same semantics across backends. Since flink seems to be the most restrictive in what's required for downstream operations consuming window_by outputs (regarding usage of a time attribute column) it makes the most sense to me to stick with the current output schemas and do the bit of extra work to massage spark's output to be semantically consistent.

[chloeh13q]

It also seems that Flink's most strict about its timestamp datatype... It requires the watermark column to be timestamp type of precision 3, and returns window_start, window_end, window_time all with timestamp type of precision 3. I don't think Spark and RW distinguish between timestamps of different precision...

In Flink, window_time casts the original timecol so that it becomes timestamp(3). After windowing TVF, the original timecol becomes a regular timestamp column and window_time becomes the new time attribute field of the returned table. I'm a little hesitant about how we want to imitate this behavior in other backends - duplicating the original timecol twice?

[chloeh13q]

I was really close to getting this to work, but it turns out that you can't group by window.start and window.end separately in Spark because it loses the timestamp column... So we can't flatten the window column before groupby. The best we can do is something like

SELECT `t1`.`window`,
       `t1`.`provinceId`, 
       AVG(`t1`.`payAmount`) AS `avg_amount` FROM 
           (SELECT `__windowed`.createTime, 
                   `__windowed`.orderId, 
                   `__windowed`.payAmount, 
                   `__windowed`.payPlatform, 
                   `__windowed`.provinceId, 
                   `__windowed`.window,
                   `__windowed`.createTime AS `window_time` FROM 
                           (SELECT `t0`.*, WINDOW(`createTime`, '60 SECOND') FROM `payment_msg` AS `t0`) AS __windowed
            ) AS `t1` 
GROUP BY 1, 2

Should we just forgo trying to unify the output schema across backends?

I think a possible alternative is to combine window_by + group by + agg into a single API, with the caveat that users cannot use the output of the window TVF as a normal table. In that case it'd be easier for us to reconcile the differences across the underlying engines.