more

Author: sryza

bin/spark-pipelines

@@ -30,4 +30,4 @@ fi
 export PYTHONPATH="${SPARK_HOME}/python/:$PYTHONPATH"
 export PYTHONPATH="${SPARK_HOME}/python/lib/py4j-0.10.9.9-src.zip:$PYTHONPATH"
 
-${SPARK_HOME}/bin/spark-submit --conf spark.api.mode=connect "${SPARK_HOME}"/python/pyspark/sql/pipelines/cli.py "$@"
+$PYSPARK_PYTHON "${SPARK_HOME}"/python/pyspark/pipelines/cli.py "$@"

dev/sparktestsupport/modules.py

@@ -411,10 +411,17 @@ def __hash__(self):
 pipelines = Module(
     name="pipelines",
     dependencies=[],
-    source_file_regexes=["sql/pipelines"],
+    source_file_regexes=["sql/pipelines", "python/pyspark/pipelines"],
     sbt_test_goals=[
         "pipelines/test",
     ],
+    python_test_goals=[
+        "pyspark.pipelines.tests.test_block_connect_access",
+        "pyspark.pipelines.tests.test_cli",
+        "pyspark.pipelines.tests.test_decorators",
+        "pyspark.pipelines.tests.test_graph_element_registry",
+        "pyspark.pipelines.tests.test_init_cli",
+    ],
 )
 
 connect = Module(
@@ -556,10 +563,6 @@ def __hash__(self):
         "pyspark.sql.tests.pandas.test_pandas_udf_window",
         "pyspark.sql.tests.pandas.test_pandas_sqlmetrics",
         "pyspark.sql.tests.pandas.test_converter",
-        "pyspark.sql.tests.pipelines.test_block_connect_access",
-        "pyspark.sql.tests.pipelines.test_cli",
-        "pyspark.sql.tests.pipelines.test_decorators",
-        "pyspark.sql.tests.pipelines.test_graph_element_registry",
         "pyspark.sql.tests.test_python_datasource",
         "pyspark.sql.tests.test_python_streaming_datasource",
         "pyspark.sql.tests.test_readwriter",

python/mypy.ini

@@ -117,6 +117,9 @@ ignore_errors = True
 [mypy-pyspark.pandas.tests.*]
 ignore_errors = True
 
+[mypy-pyspark.pipelines.tests.*]
+ignore_errors = True
+
 [mypy-pyspark.tests.*]
 ignore_errors = True
 
@@ -191,6 +194,3 @@ ignore_missing_imports = True
 ; Ignore errors for proto generated code
 [mypy-pyspark.sql.connect.proto.*, pyspark.sql.connect.proto, pyspark.sql.streaming.proto]
 ignore_errors = True
-
-[mypy-pyspark.sql.pipelines.proto.*]
-ignore_errors = True

python/pyspark/sql/pipelines/__init__.py renamed to python/pyspark/pipelines/__init__.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from pyspark.sql.pipelines.api import (
+from pyspark.pipelines.api import (
     append_flow,
     create_streaming_table,
     materialized_view,

python/pyspark/sql/pipelines/api.py renamed to python/pyspark/pipelines/api.py

@@ -17,13 +17,13 @@
 from typing import Callable, Dict, List, Optional, Union, overload
 
 from pyspark.errors import PySparkTypeError
-from pyspark.sql.pipelines.graph_element_registry import get_active_graph_element_registry
-from pyspark.sql.pipelines.type_error_utils import validate_optional_list_of_str_arg
-from pyspark.sql.pipelines.flow import Flow, QueryFunction
-from pyspark.sql.pipelines.source_code_location import (
+from pyspark.pipelines.graph_element_registry import get_active_graph_element_registry
+from pyspark.pipelines.type_error_utils import validate_optional_list_of_str_arg
+from pyspark.pipelines.flow import Flow, QueryFunction
+from pyspark.pipelines.source_code_location import (
     get_caller_source_code_location,
 )
-from pyspark.sql.pipelines.dataset import (
+from pyspark.pipelines.dataset import (
     MaterializedView,
     StreamingTable,
     TemporaryView,

python/pyspark/sql/pipelines/block_connect_access.py renamed to python/pyspark/pipelines/block_connect_access.py

@@ -32,15 +32,15 @@
 
 from pyspark.errors import PySparkException, PySparkTypeError
 from pyspark.sql import SparkSession
-from pyspark.sql.pipelines.graph_element_registry import (
+from pyspark.pipelines.graph_element_registry import (
     graph_element_registration_context,
     GraphElementRegistry,
 )
-from pyspark.sql.pipelines.init_cli import init
-from pyspark.sql.pipelines.spark_connect_graph_element_registry import (
+from pyspark.pipelines.init_cli import init
+from pyspark.pipelines.spark_connect_graph_element_registry import (
     SparkConnectGraphElementRegistry,
 )
-from pyspark.sql.pipelines.spark_connect_pipeline import (
+from pyspark.pipelines.spark_connect_pipeline import (
     create_dataflow_graph,
     start_run,
     handle_pipeline_events,

python/pyspark/sql/pipelines/cli.py renamed to python/pyspark/pipelines/cli.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass
 from typing import Mapping, Optional, Sequence, Union
 
-from pyspark.sql.pipelines.source_code_location import SourceCodeLocation
+from pyspark.pipelines.source_code_location import SourceCodeLocation
 from pyspark.sql.types import StructType
 
 

python/pyspark/sql/pipelines/dataset.py renamed to python/pyspark/pipelines/dataset.py

@@ -18,7 +18,7 @@
 from typing import Callable, Dict
 
 from pyspark.sql import DataFrame
-from pyspark.sql.pipelines.source_code_location import SourceCodeLocation
+from pyspark.pipelines.source_code_location import SourceCodeLocation
 
 QueryFunction = Callable[[], DataFrame]
 

python/pyspark/sql/pipelines/flow.py renamed to python/pyspark/pipelines/flow.py

@@ -18,8 +18,8 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 
-from pyspark.sql.pipelines.dataset import Dataset
-from pyspark.sql.pipelines.flow import Flow
+from pyspark.pipelines.dataset import Dataset
+from pyspark.pipelines.flow import Flow
 from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import Generator, Optional

python/pyspark/sql/pipelines/graph_element_registry.py renamed to python/pyspark/pipelines/graph_element_registry.py

@@ -26,8 +26,8 @@
 """
 
 PYTHON_EXAMPLE = """
+from pyspark import pipelines as sdp
 from pyspark.sql import DataFrame, SparkSession
-from pyspark.sql import pipelines as sdp
 
 spark = SparkSession.active()
 

python/pyspark/sql/pipelines/init_cli.py renamed to python/pyspark/pipelines/init_cli.py

@@ -19,25 +19,27 @@
 from pyspark.errors import PySparkTypeError
 from pyspark.sql import SparkSession
 from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
-from pyspark.sql.pipelines.block_connect_access import block_spark_connect_execution_and_analysis
-from pyspark.sql.pipelines.dataset import (
+from pyspark.pipelines.block_connect_access import block_spark_connect_execution_and_analysis
+from pyspark.pipelines.dataset import (
     Dataset,
     MaterializedView,
     Table,
     StreamingTable,
     TemporaryView,
 )
-from pyspark.sql.pipelines.flow import Flow
-from pyspark.sql.pipelines.graph_element_registry import GraphElementRegistry
-from typing import cast
+from pyspark.pipelines.flow import Flow
+from pyspark.pipelines.graph_element_registry import GraphElementRegistry
+from typing import Any, cast
 import pyspark.sql.connect.proto as pb2
 
 
 class SparkConnectGraphElementRegistry(GraphElementRegistry):
     """Registers datasets and flows in a dataflow graph held in a Spark Connect server."""
 
     def __init__(self, spark: SparkSession, dataflow_graph_id: str) -> None:
-        self._spark = spark
+        # Cast because mypy seems to think `spark`` is a function, not an object. Likely related to
+        # SPARK-47544.
+        self._client = cast(Any, spark).client
         self._dataflow_graph_id = dataflow_graph_id
 
     def register_dataset(self, dataset: Dataset) -> None:
@@ -80,12 +82,12 @@ def register_dataset(self, dataset: Dataset) -> None:
         )
         command = pb2.Command()
         command.pipeline_command.define_dataset.CopyFrom(inner_command)
-        self._spark.client.execute_command(command)
+        self._client.execute_command(command)
 
     def register_flow(self, flow: Flow) -> None:
         with block_spark_connect_execution_and_analysis():
             df = flow.func()
-        relation = cast(ConnectDataFrame, df)._plan.plan(self._spark.client)
+        relation = cast(ConnectDataFrame, df)._plan.plan(self._client)
 
         inner_command = pb2.PipelineCommand.DefineFlow(
             dataflow_graph_id=self._dataflow_graph_id,
@@ -97,7 +99,7 @@ def register_flow(self, flow: Flow) -> None:
         )
         command = pb2.Command()
         command.pipeline_command.define_flow.CopyFrom(inner_command)
-        self._spark.client.execute_command(command)
+        self._client.execute_command(command)
 
     def register_sql(self, sql_text: str, file_path: Path) -> None:
         inner_command = pb2.DefineSqlGraphElements(
@@ -107,4 +109,4 @@ def register_sql(self, sql_text: str, file_path: Path) -> None:
         )
         command = pb2.Command()
         command.pipeline_command.define_sql_graph_elements.CopyFrom(inner_command)
-        self._spark.client.execute_command(command)
+        self._client.execute_command(command)

python/pyspark/sql/pipelines/source_code_location.py renamed to python/pyspark/pipelines/source_code_location.py

@@ -14,11 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Any, Dict, Mapping, Iterator, Optional
-
-from pyspark.sql import SparkSession
+from typing import Any, Dict, Mapping, Iterator, Optional, cast
 
 import pyspark.sql.connect.proto as pb2
+from pyspark.sql import SparkSession
 from pyspark.errors.exceptions.base import PySparkValueError
 
 
@@ -39,7 +38,9 @@ def create_dataflow_graph(
     )
     command = pb2.Command()
     command.pipeline_command.create_dataflow_graph.CopyFrom(inner_command)
-    (_, properties, _) = spark.client.execute_command(command)
+    # Cast because mypy seems to think `spark`` is a function, not an object. Likely related to
+    # SPARK-47544.
+    (_, properties, _) = cast(Any, spark).client.execute_command(command)
     return properties["pipeline_command_result"].create_dataflow_graph_result.dataflow_graph_id
 
 
@@ -69,4 +70,6 @@ def start_run(spark: SparkSession, dataflow_graph_id: str) -> Iterator[Dict[str,
     inner_command = pb2.PipelineCommand.StartRun(dataflow_graph_id=dataflow_graph_id)
     command = pb2.Command()
     command.pipeline_command.start_run.CopyFrom(inner_command)
-    return spark.client.execute_command_as_iterator(command)
+    # Cast because mypy seems to think `spark`` is a function, not an object. Likely related to
+    # SPARK-47544.
+    return cast(Any, spark).client.execute_command_as_iterator(command)

python/pyspark/sql/pipelines/spark_connect_graph_element_registry.py renamed to python/pyspark/pipelines/spark_connect_graph_element_registry.py

@@ -19,9 +19,9 @@
 from pathlib import Path
 from typing import List, Sequence
 
-from pyspark.sql.pipelines.dataset import Dataset
-from pyspark.sql.pipelines.flow import Flow
-from pyspark.sql.pipelines.graph_element_registry import GraphElementRegistry
+from pyspark.pipelines.dataset import Dataset
+from pyspark.pipelines.flow import Flow
+from pyspark.pipelines.graph_element_registry import GraphElementRegistry
 
 
 @dataclass(frozen=True)

python/pyspark/sql/pipelines/spark_connect_pipeline.py renamed to python/pyspark/pipelines/spark_connect_pipeline.py

@@ -18,7 +18,7 @@
 
 from pyspark.errors import PySparkException
 from pyspark.testing.connectutils import ReusedConnectTestCase
-from pyspark.sql.pipelines.block_connect_access import block_spark_connect_execution_and_analysis
+from pyspark.pipelines.block_connect_access import block_spark_connect_execution_and_analysis
 from pyspark.testing.connectutils import (
     ReusedConnectTestCase,
     should_test_connect,

python/pyspark/sql/tests/pipelines/__init__.py renamed to python/pyspark/pipelines/tests/__init__.py

@@ -21,17 +21,16 @@
 from pathlib import Path
 
 from pyspark.errors import PySparkException
-from pyspark.sql.pipelines.cli import (
+from pyspark.pipelines.cli import (
     change_dir,
     find_pipeline_spec,
-    init,
     load_pipeline_spec,
     register_definitions,
     unpack_pipeline_spec,
     DefinitionsGlob,
     PipelineSpec,
 )
-from pyspark.sql.tests.pipelines.local_graph_element_registry import LocalGraphElementRegistry
+from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry
 
 
 class CLIUtilityTests(unittest.TestCase):
@@ -210,7 +209,7 @@ def test_register_definitions(self):
                 f.write(
                     textwrap.dedent(
                         """
-                        from pyspark.sql import pipelines as sdp
+                        from pyspark import pipelines as sdp
                         @sdp.materialized_view
                         def mv1():
                             raise NotImplementedError()
@@ -222,7 +221,7 @@ def mv1():
                 f.write(
                     textwrap.dedent(
                         """
-                        from pyspark.sql import pipelines as sdp
+                        from pyspark import pipelines as sdp
                         def mv2():
                             raise NotImplementedError()
                     """
@@ -314,27 +313,6 @@ def test_python_import_current_directory(self):
                     ),
                 )
 
-    def test_init(self):
-        with tempfile.TemporaryDirectory() as temp_dir:
-            project_name = "test_project"
-            with change_dir(Path(temp_dir)):
-                init(project_name)
-            with change_dir(Path(temp_dir) / project_name):
-                spec_path = find_pipeline_spec(Path.cwd())
-                spec = load_pipeline_spec(spec_path)
-                registry = LocalGraphElementRegistry()
-                register_definitions(spec_path, registry, spec)
-                self.assertEqual(len(registry.datasets), 1)
-                self.assertEqual(registry.datasets[0].name, "example_python_materialized_view")
-                self.assertEqual(len(registry.flows), 1)
-                self.assertEqual(registry.flows[0].name, "example_python_materialized_view")
-                self.assertEqual(registry.flows[0].target, "example_python_materialized_view")
-                self.assertEqual(len(registry.sql_files), 1)
-                self.assertEqual(
-                    registry.sql_files[0].file_path,
-                    Path("transformations") / "example_sql_materialized_view.sql",
-                )
-
 
 if __name__ == "__main__":
     try:

python/pyspark/sql/tests/pipelines/local_graph_element_registry.py renamed to python/pyspark/pipelines/tests/local_graph_element_registry.py

@@ -18,7 +18,7 @@
 import unittest
 
 from pyspark.errors import PySparkTypeError
-from pyspark.sql import pipelines as sdp
+from pyspark import pipelines as sdp
 
 
 class DecoratorsTest(unittest.TestCase):

python/pyspark/sql/tests/pipelines/test_block_connect_access.py renamed to python/pyspark/pipelines/tests/test_block_connect_access.py

@@ -18,9 +18,9 @@
 import unittest
 
 from pyspark.errors import PySparkException
-from pyspark.sql.pipelines.graph_element_registry import graph_element_registration_context
-from pyspark.sql import pipelines as sdp
-from pyspark.sql.tests.pipelines.local_graph_element_registry import LocalGraphElementRegistry
+from pyspark.pipelines.graph_element_registry import graph_element_registration_context
+from pyspark import pipelines as sdp
+from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry
 
 
 class GraphElementRegistryTest(unittest.TestCase):