using TSDF window builders

databrickslabs · Jun 2, 2024 · 7e239b5 · 7e239b5
1 parent e0af329
commit 7e239b5
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/python/tempo/interpol.py b/python/tempo/interpol.py
@@ -79,26 +79,22 @@ def interpolate(
 
     # identify transitions between segments
     seg_trans_col = "__tmp_seg_transition"
-    all_win = Window.partitionBy("symbol").orderBy("timestamp")
+    all_win = tsdf.baseWindow()
     segments = segments.withColumn(seg_trans_col,
                                    sfn.lag(needs_intpl_col, 1).over(all_win)
                                    != sfn.col(needs_intpl_col))
 
     # assign a group number to each segment
     seg_group_col = "__tmp_seg_group"
-    all_prev_win = Window.partitionBy("symbol")\
-        .orderBy("timestamp")\
-        .rowsBetween(Window.unboundedPreceding, Window.currentRow)
+    all_prev_win = tsdf.allBeforeWindow()
     segments = segments.withColumn(seg_group_col,
                                    sfn.count_if(seg_trans_col).over(all_prev_win))
 
     # build margins around intepolation segments
     if leading_margin > 0 or lagging_margin > 0:
         # collect the group number of each segment with a margin
         margin_col = "__tmp_group_with_margin"
-        margin_win = Window.partitionBy("symbol")\
-            .orderBy("timestamp")\
-            .rowsBetween(-leading_margin, lagging_margin)
+        margin_win = tsdf.rowsBetweenWindow(-leading_margin, lagging_margin)
         segments = segments.withColumn(margin_col,
                                        sfn.when(~sfn.col(needs_intpl_col),
                                                 sfn.collect_set(seg_group_col)