diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py
index c52ea62686ab9..a319dace6869a 100755
--- a/python/pyspark/ml/feature.py
+++ b/python/pyspark/ml/feature.py
@@ -5736,6 +5736,7 @@ def selectedFeatures(self):
 
 if __name__ == "__main__":
     import doctest
+    import sys
     import tempfile
 
     import pyspark.ml.feature