Address comment

Author: HyukjinKwon

python/pyspark/find_spark_home.py

@@ -36,19 +36,24 @@ def is_spark_home(path):
                 (os.path.isdir(os.path.join(path, "jars")) or
                  os.path.isdir(os.path.join(path, "assembly"))))
 
-    paths = ["../", os.path.dirname(os.path.realpath(__file__))]
+    # Spark distribution can be downloaded when HADOOP_VERSION environment variable is set.
+    # We should look up this directory first, see also SPARK-32017.
+    spark_dist_dir = "spark-distribution"
+    paths = [
+        "../",  # When we're in spark/python.
+        # Two case belows are valid when the current script is called as a library.
+        os.path.join(os.path.dirname(os.path.realpath(__file__)), spark_dist_dir),
+        os.path.dirname(os.path.realpath(__file__))]
 
     # Add the path of the PySpark module if it exists
     import_error_raised = False
     from importlib.util import find_spec
     try:
-        # Spark distribution can be downloaded when HADOOP_VERSION environment variable is set.
-        # We should look up this directory first, see also SPARK-32017.
-        spark_dist_dir = "spark-distribution"
         module_home = os.path.dirname(find_spec("pyspark").origin)
         paths.append(os.path.join(module_home, spark_dist_dir))
         paths.append(module_home)
         # If we are installed in edit mode also look two dirs up
+        # Downloading different versions are not supported in edit mode.
         paths.append(os.path.join(module_home, "../../"))
     except ImportError:
         # Not pip installed no worries