diff --git a/python/tempo/tsdf.py b/python/tempo/tsdf.py
index 6b7d462d..6a515b33 100644
--- a/python/tempo/tsdf.py
+++ b/python/tempo/tsdf.py
@@ -1362,7 +1362,7 @@ def extractStateIntervals(
         """
 
         # https://spark.apache.org/docs/latest/sql-ref-null-semantics.html#comparison-operators-
-        def null_safe_equals(col1: Column, col2: Column) -> Column[bool]:
+        def null_safe_equals(col1: Column, col2: Column) -> Column:
             return (
                 f.when(col1.isNull() & col2.isNull(), True)
                 .when(col1.isNull() | col2.isNull(), False)