TSDF: extractStateIntervals method to identify an "event_time" range where state is constant (#231)

* update gitignore

* working changes

* refactor tests & split out

* black & working changes

* comments

* working constantMetricRanges

need to complete test cases

* remove extra select to create event_ts from constantMetricRanges

* implementation and test cases for constantMetricState for PySpark comparison operator

* update constantStateRanges and new test case for state defined by column expressions

* black formatting

* update import statement for Column

* reformating and added warning message for non-tested state definitions

* move call to logger to correct spot

* fix logger warning output

* update tests and data to new format

* black formatting

* exclude cache and virtual env directories from flake8 check

* flake8 formatting

* revert formatting in examples/.

* refactor ConstantMetricState to exlude metric intervals from output

* rename method constantMetricState to extractStateIntervals

* fixing method name in warning message

Co-authored-by: Tristan Nixon <[email protected]>

Author: R7L208

.gitignore

@@ -31,7 +31,11 @@ scala/target/stream/*
 
 # ignore virtual environments
 python/venv
+python/.venv
+python/.env
 venv
+.venv
+.env
 
 # other misc ignore
 .DS_Store

python/.flake8

@@ -9,4 +9,10 @@ extend-ignore =
     # Invalid escape sequence 'x' (W605)
     W605,
     # 'from module import *' used; unable to detect undefined names (F403)
-    F403 
+    F403
+exclude =
+    __pycache__
+    env
+    .env
+    venv
+    .venv

python/tempo/tsdf.py

@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 import logging
 from functools import reduce
-from typing import List
+from typing import List, Collection, Union
 
 import numpy as np
 import pyspark.sql.functions as f
 from IPython.core.display import HTML
 from IPython.display import display as ipydisplay
 from pyspark.sql import SparkSession
+from pyspark.sql.column import Column
 from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.window import Window, WindowSpec
 from scipy.fft import fft, fftfreq
@@ -1334,6 +1337,83 @@ def tempo_fourier_util(pdf):
 
         return TSDF(result, self.ts_col, self.partitionCols, self.sequence_col)
 
+    def extractStateIntervals(
+        self,
+        *metricCols: Collection[str],
+        state_definition: Union[str, Column[bool]] = "=",
+    ) -> TSDF:
+
+        data = self.df
+
+        w = self.__baseWindow()
+
+        if type(state_definition) is str:
+            if state_definition not in ("=", "<=>", "!=", "<>", ">", "<", ">=", "<="):
+                logger.warning(
+                    "A `state_definition` which has not been tested was"
+                    "provided to the `extractStateIntervals` method."
+                )
+            current_state = f.array(*metricCols)
+        else:
+            current_state = state_definition
+
+        data = data.withColumn("current_state", current_state).drop(*metricCols)
+
+        data = (
+            data.withColumn(
+                "previous_state",
+                f.lag(f.col("current_state"), offset=1).over(w),
+            )
+            .withColumn(
+                "previous_ts",
+                f.lag(f.col(self.ts_col), offset=1).over(w),
+            )
+            .filter(f.col("previous_state").isNotNull())
+        )
+
+        if type(state_definition) is str:
+            state_change_exp = f"""
+            !(current_state {state_definition} previous_state)
+            """
+        else:
+            state_change_exp = "!(current_state AND previous_state)"
+
+        data = data.withColumn(
+            "state_change",
+            f.expr(state_change_exp),
+        ).drop("current_state", "previous_state")
+
+        data = (
+            data.withColumn(
+                "state_incrementer",
+                f.sum(f.col("state_change").cast("int")).over(w),
+            )
+            .filter(~f.col("state_change"))
+            .drop("state_change")
+        )
+
+        data = (
+            data.groupBy(*self.partitionCols, "state_incrementer")
+            .agg(
+                f.struct(
+                    f.min("previous_ts").alias("start"),
+                    f.max(f"{self.ts_col}").alias("end"),
+                ).alias(self.ts_col),
+            )
+            .drop("state_incrementer")
+        )
+
+        result = data.select(
+            self.ts_col,
+            *self.partitionCols,
+        )
+
+        return TSDF(
+            result,
+            self.ts_col,
+            self.partitionCols,
+        )
+
 
 class _ResampledTSDF(TSDF):
     def __init__(

python/tests/base.py

@@ -203,7 +203,6 @@ def assertDataFramesEqual(self, dfA, dfB):
         sortedB = dfB.select(colOrder)
         # must have identical data
         # that is all rows in A must be in B, and vice-versa
-
         self.assertEqual(sortedA.subtract(sortedB).count(), 0)
         self.assertEqual(sortedB.subtract(sortedA).count(), 0)
 

python/tests/tsdf_tests.py

@@ -3,6 +3,7 @@
 from dateutil import parser as dt_parser
 
 import pyspark.sql.functions as F
+from pyspark.sql.dataframe import DataFrame
 
 from tempo.tsdf import TSDF
 from tests.base import SparkTest
@@ -450,6 +451,140 @@ def test_upsample(self):
         self.assertDataFramesEqual(bars, barsExpected)
 
 
+class extractStateIntervalsTest(SparkTest):
+    """Test of finding time ranges for metrics with constant state."""
+
+    def create_expected_test_df(
+        self,
+        df,
+    ) -> DataFrame:
+        return (
+            # StringType not converting to TimeStamp type inside of struct so forcing
+            df.withColumn(
+                "event_ts",
+                F.struct(
+                    F.to_timestamp("event_ts.start").alias("start"),
+                    F.to_timestamp("event_ts.end").alias("end"),
+                ),
+            )
+        )
+
+    def test_eq_extractStateIntervals(self):
+
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_eq_1_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3"
+        ).df
+        extractStateIntervals_eq_2_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition="<=>"
+        ).df
+
+        # test extractStateIntervals_tsdf summary
+        self.assertDataFramesEqual(extractStateIntervals_eq_1_df, expected_df)
+        self.assertDataFramesEqual(extractStateIntervals_eq_2_df, expected_df)
+
+    def test_ne_extractStateIntervals(self):
+
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_ne_1_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition="!="
+        ).df
+        extractStateIntervals_ne_2_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition="<>"
+        ).df
+
+        # test extractStateIntervals_tsdf summary
+        self.assertDataFramesEqual(extractStateIntervals_ne_1_df, expected_df)
+        self.assertDataFramesEqual(extractStateIntervals_ne_2_df, expected_df)
+
+    def test_gt_extractStateIntervals(self):
+
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_gt_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition=">"
+        ).df
+
+        self.assertDataFramesEqual(extractStateIntervals_gt_df, expected_df)
+
+    def test_lt_extractStateIntervals(self):
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_lt_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition="<"
+        ).df
+
+        # test extractStateIntervals_tsdf summary
+        self.assertDataFramesEqual(extractStateIntervals_lt_df, expected_df)
+
+    def test_gte_extractStateIntervals(self):
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_gt_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition=">="
+        ).df
+
+        self.assertDataFramesEqual(extractStateIntervals_gt_df, expected_df)
+
+    def test_lte_extractStateIntervals(self):
+
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_lte_df = input_tsdf.extractStateIntervals(
+            "metric_1", "metric_2", "metric_3", state_definition="<="
+        ).df
+
+        # test extractStateIntervals_tsdf summary
+        self.assertDataFramesEqual(extractStateIntervals_lte_df, expected_df)
+
+    def test_bool_col_extractStateIntervals(self):
+
+        # construct dataframes
+        input_tsdf = self.get_data_as_tsdf("input")
+        expected_df = self.get_data_as_sdf("expected")
+        expected_df = self.create_expected_test_df(expected_df)
+
+        # call extractStateIntervals method
+        extractStateIntervals_bool_col_df = input_tsdf.extractStateIntervals(
+            "metric_1",
+            "metric_2",
+            "metric_3",
+            state_definition=F.abs(
+                F.col("metric_1") - F.col("metric_2") - F.col("metric_3")
+            )
+            < F.lit(10),
+        ).df
+
+        # test extractStateIntervals_tsdf summary
+        self.assertDataFramesEqual(extractStateIntervals_bool_col_df, expected_df)
+
+
 # MAIN
 if __name__ == "__main__":
     unittest.main()