Skip to content

Commit 7e239b5

Browse files
author
Tristan Nixon
committed
using TSDF window builders
1 parent e0af329 commit 7e239b5

File tree

1 file changed

+3
-7
lines changed

1 file changed

+3
-7
lines changed

python/tempo/interpol.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,26 +79,22 @@ def interpolate(
7979

8080
# identify transitions between segments
8181
seg_trans_col = "__tmp_seg_transition"
82-
all_win = Window.partitionBy("symbol").orderBy("timestamp")
82+
all_win = tsdf.baseWindow()
8383
segments = segments.withColumn(seg_trans_col,
8484
sfn.lag(needs_intpl_col, 1).over(all_win)
8585
!= sfn.col(needs_intpl_col))
8686

8787
# assign a group number to each segment
8888
seg_group_col = "__tmp_seg_group"
89-
all_prev_win = Window.partitionBy("symbol")\
90-
.orderBy("timestamp")\
91-
.rowsBetween(Window.unboundedPreceding, Window.currentRow)
89+
all_prev_win = tsdf.allBeforeWindow()
9290
segments = segments.withColumn(seg_group_col,
9391
sfn.count_if(seg_trans_col).over(all_prev_win))
9492

9593
# build margins around intepolation segments
9694
if leading_margin > 0 or lagging_margin > 0:
9795
# collect the group number of each segment with a margin
9896
margin_col = "__tmp_group_with_margin"
99-
margin_win = Window.partitionBy("symbol")\
100-
.orderBy("timestamp")\
101-
.rowsBetween(-leading_margin, lagging_margin)
97+
margin_win = tsdf.rowsBetweenWindow(-leading_margin, lagging_margin)
10298
segments = segments.withColumn(margin_col,
10399
sfn.when(~sfn.col(needs_intpl_col),
104100
sfn.collect_set(seg_group_col)

0 commit comments

Comments
 (0)