python/pyspark/sql/connect/session.py - master - githubmirror/spark

python/pyspark/sql/connect/session.py
1 443 строки55 KB
Tian Gao
[SPARK-58159][PYTHON] Support `with` statement for connect session
30 июл 2026, 03:49
30 июл 2026, 03:496c37fe3
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import uuid

import json
import threading
import os
import sys
import warnings
from collections.abc import Callable, Sized
import functools
from threading import RLock
from types import TracebackType
from typing import (
    Optional,
    Any,
    Iterator,
    Union,
    Dict,
    List,
    Tuple,
    Type,
    Set,
    cast,
    overload,
    Iterable,
    Mapping,
    TYPE_CHECKING,
    ClassVar,
)

import numpy as np
import pandas as pd
import pyarrow as pa
from pandas.api.types import is_datetime64_dtype, is_timedelta64_dtype
import urllib

from pyspark.sql.connect.dataframe import DataFrame
from pyspark.sql.dataframe import DataFrame as ParentDataFrame
from pyspark.sql.connect.logging import logger
from pyspark.sql.connect.client import SparkConnectClient, DefaultChannelBuilder
from pyspark.sql.connect.conf import RuntimeConf
from pyspark.sql.connect.plan import (
    SQL,
    Range,
    LocalRelation,
    LogicalPlan,
    ChunkedCachedLocalRelation,
    CachedRelation,
    CachedRemoteRelation,
    SubqueryAlias,
)
from pyspark.sql.connect.functions import builtin as F
from pyspark.sql.connect.profiler import ProfilerCollector
from pyspark.sql.connect.readwriter import DataFrameReader
from pyspark.sql.connect.streaming.readwriter import DataStreamReader
from pyspark.sql.connect.streaming.query import StreamingQueryManager
from pyspark.sql.pandas.conversion import create_arrow_table_from_pandas
from pyspark.sql.pandas.types import (
    to_arrow_schema,
    _deduplicate_field_names,
    from_arrow_schema,
    from_arrow_type,
    _check_arrow_table_timestamps_localize,
)
from pyspark.sql.profiler import Profile
from pyspark.sql.session import classproperty, SparkSession as PySparkSession
from pyspark.sql.types import (
    _infer_schema,
    _has_nulltype,
    _merge_type,
    Row,
    DataType,
    DayTimeIntervalType,
    StructType,
    AtomicType,
    TimestampType,
    MapType,
    StringType,
)
from pyspark.sql.utils import to_str
from pyspark.errors import (
    AnalysisException,
    PySparkAttributeError,
    PySparkNotImplementedError,
    PySparkRuntimeError,
    PySparkValueError,
    PySparkTypeError,
    PySparkAssertionError,
)

if TYPE_CHECKING:
    import pyspark.sql.connect.proto as pb2
    from pyspark.sql.connect._typing import OptionalPrimitiveType
    from pyspark.sql.connect.catalog import Catalog
    from pyspark.sql.connect.udf import UDFRegistration
    from pyspark.sql.connect.udtf import UDTFRegistration
    from pyspark.sql.connect.tvf import TableValuedFunction
    from pyspark.sql.connect.shell.progress import ProgressHandler
    from pyspark.sql.connect.datasource import DataSourceRegistration


class SparkSession:
    # The active SparkSession for the current thread
    _active_session: ClassVar[threading.local] = threading.local()
    # Reference to the root SparkSession
    _default_session: ClassVar[Optional["SparkSession"]] = None
    _lock: ClassVar[RLock] = RLock()

    class Builder:
        """Builder for :class:`SparkSession`."""

        _lock = RLock()

        def __init__(self) -> None:
            self._options: Dict[str, Any] = {}
            self._channel_builder: Optional[DefaultChannelBuilder] = None
            self._hook_factories: list["Callable[[SparkSession], SparkSession.Hook]"] = []

        @overload
        def config(self, key: str, value: Any) -> "SparkSession.Builder": ...

        @overload
        def config(self, *, map: Dict[str, "OptionalPrimitiveType"]) -> "SparkSession.Builder": ...

        def config(
            self,
            key: Optional[str] = None,
            value: Optional[Any] = None,
            *,
            map: Optional[Dict[str, "OptionalPrimitiveType"]] = None,
        ) -> "SparkSession.Builder":
            with self._lock:
                if map is not None:
                    for k, v in map.items():
                        self._options[k] = to_str(v)
                else:
                    self._options[cast(str, key)] = to_str(value)
                return self

        def master(self, master: str) -> "SparkSession.Builder":
            return self

        def appName(self, name: str) -> "SparkSession.Builder":
            return self.config("spark.app.name", name)

        def remote(self, location: str = "sc://localhost") -> "SparkSession.Builder":
            return self.config("spark.remote", location)

        def channelBuilder(self, channelBuilder: DefaultChannelBuilder) -> "SparkSession.Builder":
            """Uses custom :class:`ChannelBuilder` implementation, when there is a need
            to customize the behavior for creation of GRPC connections.

            .. versionadded:: 3.5.0

            An example to use this class looks like this:

            .. code-block:: python

                from pyspark.sql.connect import SparkSession, ChannelBuilder

                class CustomChannelBuilder(ChannelBuilder):
                    ...

                custom_channel_builder = CustomChannelBuilder(...)
                spark = SparkSession.builder().channelBuilder(custom_channel_builder).getOrCreate()

            Returns
            -------
            :class:`SparkSession.Builder`
            """
            with self._lock:
                # self._channel_builder is a separate field, because it may hold the state
                # and cannot be serialized with to_str()
                self._channel_builder = channelBuilder
                return self

        def _registerHook(
            self, hook_factory: "Callable[[SparkSession], SparkSession.Hook]"
        ) -> "SparkSession.Builder":
            with self._lock:
                self._hook_factories.append(hook_factory)
                return self

        def enableHiveSupport(self) -> "SparkSession.Builder":
            raise PySparkNotImplementedError(
                errorClass="NOT_IMPLEMENTED", messageParameters={"feature": "enableHiveSupport"}
            )

        def _apply_options(self, session: "SparkSession") -> None:
            init_opts = {}
            for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))):
                init_opts = json.loads(os.environ[f"PYSPARK_REMOTE_INIT_CONF_{i}"])

            # The options are applied after session creation,
            # so options ["spark.remote", "spark.master"] always take no effect.
            invalid_opts = ["spark.remote", "spark.master"]

            with self._lock:
                opts = {}

                # Only attempts to set Spark SQL configurations.
                # If the configurations are static, it might throw an exception so
                # simply ignore it for now.
                for k, v in init_opts.items():
                    if k not in invalid_opts and k.startswith("spark.sql."):
                        opts[k] = v

                for k, v in self._options.items():
                    if k not in invalid_opts:
                        opts[k] = v

                if len(opts) > 0:
                    session.conf._set_all(configs=opts, silent=True)

        def create(self) -> "SparkSession":
            has_channel_builder = self._channel_builder is not None
            has_spark_remote = "spark.remote" in self._options

            if (has_channel_builder and has_spark_remote) or (
                not has_channel_builder and not has_spark_remote
            ):
                raise PySparkValueError(
                    errorClass="SESSION_NEED_CONN_STR_OR_BUILDER", messageParameters={}
                )

            if has_channel_builder:
                assert self._channel_builder is not None
                session = SparkSession(
                    connection=self._channel_builder, hook_factories=self._hook_factories
                )
            else:
                spark_remote = to_str(self._options.get("spark.remote"))
                assert spark_remote is not None
                session = SparkSession(connection=spark_remote, hook_factories=self._hook_factories)

            SparkSession._set_default_and_active_session(session)
            self._apply_options(session)
            return session

        def getOrCreate(self) -> "SparkSession":
            with SparkSession._lock:
                session = SparkSession.getActiveSession()
                if session is None:
                    session = SparkSession._get_default_session()
                    if session is None:
                        session = self.create()
                self._apply_options(session)
                return session

    class Hook:
        """A Hook can be used to inject behavior into the session."""

        def on_execute_plan(self, request: "pb2.ExecutePlanRequest") -> "pb2.ExecutePlanRequest":
            """Called before sending an ExecutePlanRequest.

            The request is replaced with the one returned by this method.
            """
            return request

    _client: SparkConnectClient

    # SPARK-47544: Explicitly declaring this as an identifier instead of a method.
    # If changing, make sure this bug is not reintroduced.
    builder: Builder = classproperty(lambda cls: cls.Builder())  # type: ignore
    builder.__doc__ = PySparkSession.builder.__doc__

    def __init__(
        self,
        connection: Union[str, DefaultChannelBuilder],
        userId: Optional[str] = None,
        hook_factories: Optional[list["Callable[[SparkSession], Hook]"]] = None,
    ) -> None:
        """
        Creates a new SparkSession for the Spark Connect interface.

        Parameters
        ----------
        connection: str or class:`ChannelBuilder`
            Connection string that is used to extract the connection parameters and configure
            the GRPC connection. Or instance of ChannelBuilder that creates GRPC connection.
            Defaults to `sc://localhost`.
        userId : str, optional
            Optional unique user ID that is used to differentiate multiple users and
            isolate their Spark Sessions. If the `user_id` is not set, will default to
            the $USER environment. Defining the user ID as part of the connection string
            takes precedence.
        hook_factories: list[Callable[[SparkSession], Hook]], optional
            Optional list of hook factories for hooks that should be registered for this session.
        """
        hook_factories = hook_factories or []
        self._client = SparkConnectClient(
            connection=connection,
            user_id=userId,
            session_hooks=[factory(self) for factory in hook_factories],
        )
        self._session_id = self._client._session_id

        # Set to false to prevent client.release_session on close() (testing only)
        self.release_session_on_close = True

    @classmethod
    def _set_default_and_active_session(cls, session: "SparkSession") -> None:
        """
        Set the (global) default :class:`SparkSession`, and (thread-local)
        active :class:`SparkSession` when they are not set yet.
        """
        with cls._lock:
            if cls._default_session is None:
                cls._default_session = session
        if getattr(cls._active_session, "session", None) is None:
            cls._active_session.session = session

    @classmethod
    def _get_default_session(cls) -> Optional["SparkSession"]:
        s = cls._default_session
        if s is not None and not s.is_stopped:
            return s
        return None

    @classmethod
    def getActiveSession(cls) -> Optional["SparkSession"]:
        s = getattr(cls._active_session, "session", None)
        if s is not None and not s.is_stopped:
            return s
        return None

    @classmethod
    def _getActiveSessionIfMatches(cls, session_id: str) -> "SparkSession":
        """
        Internal use only. This method is called from the custom handler
        generated by __reduce__. To avoid serializing a WeakRef, we create a
        custom classmethod to instantiate the SparkSession.
        """
        session = SparkSession.getActiveSession()
        if session is None:
            raise PySparkRuntimeError(
                errorClass="NO_ACTIVE_SESSION",
                messageParameters={},
            )
        if session._session_id != session_id:
            raise PySparkAssertionError(
                "Expected session ID does not match active session ID: "
                f"{session_id} != {session._session_id}"
            )
        return session

    getActiveSession.__doc__ = PySparkSession.getActiveSession.__doc__

    @classmethod
    def active(cls) -> "SparkSession":
        session = cls.getActiveSession()
        if session is None:
            session = cls._get_default_session()
            if session is None:
                raise PySparkRuntimeError(
                    errorClass="NO_ACTIVE_OR_DEFAULT_SESSION",
                    messageParameters={},
                )
        return session

    active.__doc__ = PySparkSession.active.__doc__

    def table(self, tableName: str) -> ParentDataFrame:
        if not isinstance(tableName, str):
            raise PySparkTypeError(
                errorClass="NOT_EXPECTED_TYPE",
                messageParameters={
                    "arg_name": "tableName",
                    "expected_type": "str",
                    "arg_type": type(tableName).__name__,
                },
            )

        return self.read.table(tableName)

    table.__doc__ = PySparkSession.table.__doc__

    @property
    def read(self) -> "DataFrameReader":
        return DataFrameReader(self)

    read.__doc__ = PySparkSession.read.__doc__

    @property
    def readStream(self) -> "DataStreamReader":
        return DataStreamReader(self)

    readStream.__doc__ = PySparkSession.readStream.__doc__

    @property
    def tvf(self) -> "TableValuedFunction":
        from pyspark.sql.connect.tvf import TableValuedFunction

        return TableValuedFunction(self)

    tvf.__doc__ = PySparkSession.tvf.__doc__

    def registerProgressHandler(self, handler: "ProgressHandler") -> None:
        self._client.register_progress_handler(handler)

    registerProgressHandler.__doc__ = PySparkSession.registerProgressHandler.__doc__

    def removeProgressHandler(self, handler: "ProgressHandler") -> None:
        self._client.remove_progress_handler(handler)

    removeProgressHandler.__doc__ = PySparkSession.removeProgressHandler.__doc__

    def clearProgressHandlers(self) -> None:
        self._client.clear_progress_handlers()

    clearProgressHandlers.__doc__ = PySparkSession.clearProgressHandlers.__doc__

    def _inferSchemaFromList(
        self,
        data: Iterable[Any],
        names: Optional[List[str]],
        configs: Mapping[str, Optional[str]],
    ) -> StructType:
        """
        Infer schema from list of Row, dict, or tuple.
        """
        if not data:
            raise PySparkValueError(
                errorClass="CANNOT_INFER_EMPTY_SCHEMA",
                messageParameters={},
            )

        (
            infer_dict_as_struct,
            infer_array_from_first_element,
            infer_map_from_first_pair,
            prefer_timestamp,
        ) = (
            configs["spark.sql.pyspark.inferNestedDictAsStruct.enabled"],
            configs["spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled"],
            configs["spark.sql.pyspark.legacy.inferMapTypeFromFirstPair.enabled"],
            configs["spark.sql.timestampType"],
        )
        return functools.reduce(
            _merge_type,
            (
                _infer_schema(
                    row,
                    names,
                    infer_dict_as_struct=(infer_dict_as_struct == "true"),
                    infer_array_from_first_element=(infer_array_from_first_element == "true"),
                    infer_map_from_first_pair=(infer_map_from_first_pair == "true"),
                    prefer_timestamp_ntz=(prefer_timestamp == "TIMESTAMP_NTZ"),
                )
                for row in data
            ),
        )

    def createDataFrame(
        self,
        data: Union["pd.DataFrame", "np.ndarray", "pa.Table", Iterable[Any]],
        schema: Optional[Union[AtomicType, StructType, str, List[str], Tuple[str, ...]]] = None,
        samplingRatio: Optional[float] = None,
        verifySchema: Optional[bool] = None,
    ) -> "ParentDataFrame":
        assert data is not None
        if isinstance(data, DataFrame):
            raise PySparkTypeError(
                errorClass="INVALID_TYPE",
                messageParameters={"arg_name": "data", "arg_type": "DataFrame"},
            )

        if samplingRatio is not None:
            warnings.warn("'samplingRatio' is ignored. It is not supported with Spark Connect.")

        if verifySchema is not None:
            warnings.warn("'verifySchema' is ignored. It is not supported with Spark Connect.")

        _schema: Optional[Union[AtomicType, StructType]] = None
        _cols: Optional[List[str]] = None
        _num_cols: Optional[int] = None

        if isinstance(schema, str):
            schema = self.client._analyze(  # type: ignore[assignment]
                method="ddl_parse", ddl_string=schema
            ).parsed

        if isinstance(schema, (AtomicType, StructType)):
            _schema = schema
            if isinstance(schema, StructType):
                _num_cols = len(schema.fields)
            else:
                _num_cols = 1

        elif isinstance(schema, (list, tuple)):
            # Must re-encode any unicode strings to be consistent with StructField names
            _cols = [x.encode("utf-8") if not isinstance(x, str) else x for x in schema]
            _num_cols = len(_cols)

        elif schema is not None:
            raise PySparkTypeError(
                errorClass="NOT_EXPECTED_TYPE",
                messageParameters={
                    "arg_name": "schema",
                    "expected_type": "list, None or StructType",
                    "arg_type": type(schema).__name__,
                },
            )

        if isinstance(data, np.ndarray) and data.ndim not in [1, 2]:
            raise PySparkValueError(
                errorClass="INVALID_NDARRAY_DIMENSION",
                messageParameters={"dimensions": "1 or 2"},
            )
        elif isinstance(data, Sized) and len(data) == 0:
            if _schema is not None:
                return DataFrame(LocalRelation(table=None, schema=_schema.json()), self)
            else:
                raise PySparkValueError(
                    errorClass="CANNOT_INFER_EMPTY_SCHEMA",
                    messageParameters={},
                )

        # Get all related configs in a batch
        configs = self._client.get_config_dict(
            "spark.sql.timestampType",
            "spark.sql.session.timeZone",
            "spark.sql.session.localRelationCacheThreshold",
            "spark.sql.session.localRelationSizeLimit",
            "spark.sql.session.localRelationChunkSizeRows",
            "spark.sql.session.localRelationChunkSizeBytes",
            "spark.sql.session.localRelationBatchOfChunksSizeBytes",
            "spark.sql.execution.pandas.convertToArrowArraySafely",
            "spark.sql.execution.pandas.inferPandasDictAsMap",
            "spark.sql.pyspark.inferNestedDictAsStruct.enabled",
            "spark.sql.pyspark.legacy.inferArrayTypeFromFirstElement.enabled",
            "spark.sql.pyspark.legacy.inferMapTypeFromFirstPair.enabled",
            "spark.sql.execution.arrow.useLargeVarTypes",
        )
        timezone = configs["spark.sql.session.timeZone"]
        prefer_timestamp_ntz = configs["spark.sql.timestampType"] == "TIMESTAMP_NTZ"
        prefers_large_types = configs["spark.sql.execution.arrow.useLargeVarTypes"] == "true"

        _table: Optional[pa.Table] = None

        if isinstance(data, pd.DataFrame):
            # Logic was borrowed from `_create_from_pandas_with_arrow` in
            # `pyspark.sql.pandas.conversion.py`. Should ideally deduplicate the logics.

            # If no schema supplied by user then get the names of columns only
            if schema is None:
                _cols = [str(x) if not isinstance(x, str) else x for x in data.columns]
                if len(_cols) == 0:
                    raise PySparkValueError(
                        errorClass="CANNOT_INFER_EMPTY_SCHEMA",
                        messageParameters={},
                    )
                infer_pandas_dict_as_map = (
                    configs["spark.sql.execution.pandas.inferPandasDictAsMap"] == "true"
                )
                if infer_pandas_dict_as_map:
                    struct = StructType()
                    pa_schema = pa.Schema.from_pandas(data)
                    spark_type: Union[MapType, DataType]
                    for field in pa_schema:
                        field_type = field.type
                        if isinstance(field_type, pa.StructType):
                            if len(field_type) == 0:
                                raise PySparkValueError(
                                    errorClass="CANNOT_INFER_EMPTY_SCHEMA",
                                    messageParameters={},
                                )
                            arrow_type = field_type.field(0).type
                            spark_type = MapType(
                                StringType(),
                                from_arrow_type(arrow_type, prefer_timestamp_ntz),
                            )
                        else:
                            spark_type = from_arrow_type(field_type, prefer_timestamp_ntz)
                        struct.add(field.name, spark_type, nullable=field.nullable)
                    schema = struct
            elif isinstance(schema, (list, tuple)) and cast(int, _num_cols) < len(data.columns):
                assert isinstance(_cols, list)
                _cols.extend([f"_{i + 1}" for i in range(cast(int, _num_cols), len(data.columns))])
                _num_cols = len(_cols)

            # Determine arrow types to coerce data when creating batches
            arrow_schema: Optional[pa.Schema] = None
            spark_types: List[Optional[DataType]]
            if isinstance(schema, StructType):
                deduped_schema = cast(StructType, _deduplicate_field_names(schema))
                spark_types = [field.dataType for field in deduped_schema.fields]
                arrow_schema = to_arrow_schema(
                    deduped_schema,
                    timezone="UTC",
                    prefers_large_types=prefers_large_types,
                )
                _cols = [str(x) if not isinstance(x, str) else x for x in schema.fieldNames()]
            elif isinstance(schema, DataType):
                raise PySparkTypeError(
                    errorClass="UNSUPPORTED_DATA_TYPE_FOR_ARROW",
                    messageParameters={"data_type": str(schema)},
                )
            else:
                # Any timestamps must be coerced to be compatible with Spark
                spark_types = [
                    (
                        TimestampType()
                        if is_datetime64_dtype(t) or isinstance(t, pd.DatetimeTZDtype)
                        else DayTimeIntervalType()
                        if is_timedelta64_dtype(t)
                        else None
                    )
                    for t in data.dtypes
                ]

            safecheck = configs["spark.sql.execution.pandas.convertToArrowArraySafely"]

            # Handle the 0-column case separately to preserve row count.
            # pa.RecordBatch.from_pandas preserves num_rows via pandas index metadata.
            if len(data.columns) == 0:
                _table = pa.Table.from_batches([pa.RecordBatch.from_pandas(data)])
            else:
                _table = create_arrow_table_from_pandas(
                    [(c, st) for (_, c), st in zip(data.items(), spark_types)],
                    timezone=cast(str, timezone),
                    safecheck=safecheck == "true",
                    prefers_large_types=prefers_large_types,
                )

            if isinstance(schema, StructType):
                assert arrow_schema is not None
                # Skip cast for 0-column tables as it loses row count
                if len(schema.fields) > 0:
                    _table = _table.rename_columns(
                        cast(StructType, _deduplicate_field_names(schema)).names
                    ).cast(arrow_schema)

        elif isinstance(data, pa.Table):
            # If no schema supplied by user then get the names of columns only
            if schema is None:
                _cols = data.column_names
            if isinstance(schema, (list, tuple)) and cast(int, _num_cols) < len(data.columns):
                assert isinstance(_cols, list)
                _cols.extend([f"_{i + 1}" for i in range(cast(int, _num_cols), len(data.columns))])
                _num_cols = len(_cols)

            if not isinstance(schema, StructType):
                schema = from_arrow_schema(data.schema, prefer_timestamp_ntz=prefer_timestamp_ntz)

            _table = (
                _check_arrow_table_timestamps_localize(data, schema, True, timezone)
                .cast(
                    to_arrow_schema(
                        schema,
                        error_on_duplicated_field_names_in_struct=True,
                        timezone="UTC",
                        prefers_large_types=prefers_large_types,
                    )
                )
                .rename_columns(schema.names)
            )

        elif isinstance(data, np.ndarray):
            if _cols is None:
                if data.ndim == 1 or data.shape[1] == 1:
                    _cols = ["value"]
                else:
                    _cols = ["_%s" % i for i in range(1, data.shape[1] + 1)]

            if data.ndim == 1:
                if 1 != len(_cols):
                    raise PySparkValueError(
                        errorClass="AXIS_LENGTH_MISMATCH",
                        messageParameters={
                            "expected_length": str(len(_cols)),
                            "actual_length": "1",
                        },
                    )

                _table = pa.Table.from_arrays([pa.array(data)], _cols)
            else:
                if data.shape[1] != len(_cols):
                    raise PySparkValueError(
                        errorClass="AXIS_LENGTH_MISMATCH",
                        messageParameters={
                            "expected_length": str(len(_cols)),
                            "actual_length": str(data.shape[1]),
                        },
                    )

                _table = pa.Table.from_arrays(
                    [pa.array(data[::, i]) for i in range(0, data.shape[1])], _cols
                )

            # The _table should already have the proper column names.
            _cols = None

        else:
            _data = list(data)

            if isinstance(_data[0], dict):
                # Sort the data to respect inferred schema.
                # For dictionaries, we sort the schema in alphabetical order.
                _data = [dict(sorted(d.items())) if d is not None else None for d in _data]

            elif not isinstance(_data[0], (Row, tuple, list, dict)) and not hasattr(
                _data[0], "__dict__"
            ):
                # input data can be [1, 2, 3]
                # we need to convert it to [[1], [2], [3]] to be able to infer schema.
                _data = [[d] for d in _data]

            if _schema is not None:
                if not isinstance(_schema, StructType):
                    _schema = StructType().add("value", _schema)
            else:
                _schema = self._inferSchemaFromList(_data, _cols, configs)

                if _cols is not None and cast(int, _num_cols) < len(_cols):
                    _num_cols = len(_cols)

                if _has_nulltype(_schema):
                    # For cases like createDataFrame([("Alice", None, 80.1)], schema)
                    # we can not infer the schema from the data itself.
                    raise PySparkValueError(
                        errorClass="CANNOT_DETERMINE_TYPE", messageParameters={}
                    )

            from pyspark.sql.conversion import (
                LocalDataToArrowConversion,
            )

            # Spark Connect will try its best to build the Arrow table with the
            # inferred schema in the client side, and then rename the columns and
            # cast the datatypes in the server side.
            _table = LocalDataToArrowConversion.convert(_data, _schema, prefers_large_types)

        # TODO: Beside the validation on number of columns, we should also check
        # whether the Arrow Schema is compatible with the user provided Schema.
        if _num_cols is not None and _num_cols != _table.shape[1]:
            raise PySparkValueError(
                errorClass="AXIS_LENGTH_MISMATCH",
                messageParameters={
                    "expected_length": str(_num_cols),
                    "actual_length": str(_table.shape[1]),
                },
            )

        if _schema is not None:
            local_relation = LocalRelation(_table, schema=_schema.json())
        else:
            local_relation = LocalRelation(_table)

        # get_config_dict throws [SQL_CONF_NOT_FOUND] if the key is not found.
        cache_threshold = int(
            configs["spark.sql.session.localRelationCacheThreshold"]  # type: ignore[arg-type]
        )
        local_relation_size_limit = int(
            configs["spark.sql.session.localRelationSizeLimit"]  # type: ignore[arg-type]
        )
        max_chunk_size_rows = int(
            configs["spark.sql.session.localRelationChunkSizeRows"]  # type: ignore[arg-type]
        )
        max_chunk_size_bytes = int(
            configs["spark.sql.session.localRelationChunkSizeBytes"]  # type: ignore[arg-type]
        )
        max_batch_of_chunks_size_bytes = int(
            configs["spark.sql.session.localRelationBatchOfChunksSizeBytes"]  # type: ignore[arg-type]
        )
        plan: LogicalPlan = local_relation
        if cache_threshold <= _table.nbytes:
            plan = self._cache_local_relation(
                local_relation,
                local_relation_size_limit,
                max_chunk_size_rows,
                max_chunk_size_bytes,
                max_batch_of_chunks_size_bytes,
            )

        df = DataFrame(plan, self)
        if _cols is not None and len(_cols) > 0:
            df = df.toDF(*_cols)  # type: ignore[assignment]
        return df

    createDataFrame.__doc__ = PySparkSession.createDataFrame.__doc__

    def emptyDataFrame(self, schema: Union[StructType, str]) -> "ParentDataFrame":
        return self.createDataFrame([], schema)

    emptyDataFrame.__doc__ = PySparkSession.emptyDataFrame.__doc__

    def sql(
        self,
        sqlQuery: str,
        args: Optional[Union[Dict[str, Any], List[Any]]] = None,
        **kwargs: Any,
    ) -> "ParentDataFrame":
        _args = []
        _named_args = {}
        if args is not None:
            if isinstance(args, Dict):
                for k, v in args.items():
                    assert isinstance(k, str)
                    _named_args[k] = F.lit(v)
            elif isinstance(args, List):
                _args = [F.lit(v) for v in args]
            else:
                raise PySparkTypeError(
                    errorClass="INVALID_TYPE",
                    messageParameters={"arg_name": "args", "arg_type": type(args).__name__},
                )

        _views: List[SubqueryAlias] = []
        if len(kwargs) > 0:
            from pyspark.sql.connect.sql_formatter import SQLStringFormatter

            formatter = SQLStringFormatter(self)
            sqlQuery = formatter.format(sqlQuery, **kwargs)

            for df, name in formatter._temp_views:
                _views.append(SubqueryAlias(df._plan, name))

        cmd = SQL(sqlQuery, _args, _named_args, _views)
        data, properties, ei = self.client.execute_command(cmd.command(self._client))
        if "sql_command_result" in properties:
            df = DataFrame(CachedRelation(properties["sql_command_result"]), self)
            # A command result contains the execution.
            df._execution_info = ei
            return df
        else:
            return DataFrame(cmd, self)

    sql.__doc__ = PySparkSession.sql.__doc__

    def range(
        self,
        start: int,
        end: Optional[int] = None,
        step: int = 1,
        numPartitions: Optional[int] = None,
    ) -> ParentDataFrame:
        if end is None:
            actual_end = start
            start = 0
        else:
            actual_end = end

        if numPartitions is not None:
            numPartitions = int(numPartitions)

        return DataFrame(
            Range(
                start=int(start), end=int(actual_end), step=int(step), num_partitions=numPartitions
            ),
            self,
        )

    range.__doc__ = PySparkSession.range.__doc__

    @functools.cached_property
    def catalog(self) -> "Catalog":
        from pyspark.sql.connect.catalog import Catalog

        return Catalog(self)

    catalog.__doc__ = PySparkSession.catalog.__doc__

    def __del__(self) -> None:
        try:
            # StreamingQueryManager has client states that needs to be cleaned up
            if hasattr(self, "_sqm"):
                self._sqm.close()
            # Try its best to close.
            self.client.close()
        except Exception:
            pass

    def interruptAll(self) -> List[str]:
        op_ids = self.client.interrupt_all()
        assert op_ids is not None
        return op_ids

    interruptAll.__doc__ = PySparkSession.interruptAll.__doc__

    def interruptTag(self, tag: str) -> List[str]:
        op_ids = self.client.interrupt_tag(tag)
        assert op_ids is not None
        return op_ids

    interruptTag.__doc__ = PySparkSession.interruptTag.__doc__

    def interruptOperation(self, op_id: str) -> List[str]:
        op_ids = self.client.interrupt_operation(op_id)
        assert op_ids is not None
        return op_ids

    interruptOperation.__doc__ = PySparkSession.interruptOperation.__doc__

    def addTag(self, tag: str) -> None:
        self.client.add_tag(tag)

    addTag.__doc__ = PySparkSession.addTag.__doc__

    def removeTag(self, tag: str) -> None:
        self.client.remove_tag(tag)

    removeTag.__doc__ = PySparkSession.removeTag.__doc__

    def getTags(self) -> Set[str]:
        return self.client.get_tags()

    getTags.__doc__ = PySparkSession.getTags.__doc__

    def clearTags(self) -> None:
        return self.client.clear_tags()

    clearTags.__doc__ = PySparkSession.clearTags.__doc__

    def stop(self) -> None:
        """
        Release the current session and close the GRPC connection to the Spark Connect server.
        Reset the active session so that calls to getOrCreate() creates a new session.
        If the session was created in local mode, the Spark Connect server running locally is also
        terminated.

        This API is best-effort and idempotent, i.e., if any of the operations fail, the API will
        not produce an error. Stopping an already stopped session is a no-op.
        """

        # Whereas the regular PySpark session immediately terminates the Spark Context
        # itself, meaning that stopping all Spark sessions, this will only stop this one session
        # on the server.
        # It is controversial to follow the existing the regular Spark session's behavior
        # specifically in Spark Connect the Spark Connect server is designed for
        # multi-tenancy - the remote client side cannot just stop the server and stop
        # other remote clients being used from other users.
        with SparkSession._lock:
            if not self.is_stopped and self.release_session_on_close:
                try:
                    self.client.release_session()
                except Exception as e:
                    logger.warning(f"session.stop(): Session could not be released. Error: ${e}")

            try:
                self.client.close()
            except Exception as e:
                logger.warning(f"session.stop(): Client could not be closed. Error: ${e}")

            if self is SparkSession._default_session:
                SparkSession._default_session = None
            if self is getattr(SparkSession._active_session, "session", None):
                SparkSession._active_session.session = None

            # Only touch the SQLContext cache if the module was ever imported; if no
            # SQLContext was created there is nothing to reset, and we avoid importing it.
            _connect_context = sys.modules.get("pyspark.sql.connect.context")
            if _connect_context is not None:
                _ConnectSQLContext = _connect_context.SQLContext
                if (
                    _ConnectSQLContext._instantiatedContext is not None
                    and _ConnectSQLContext._instantiatedContext.sparkSession is self
                ):
                    _ConnectSQLContext._instantiatedContext = None

            if "SPARK_LOCAL_REMOTE" in os.environ:
                # When local mode is in use, follow the regular Spark session's
                # behavior by terminating the Spark Connect server,
                # meaning that you can stop local mode, and restart the Spark Connect
                # client with a different remote address.
                if PySparkSession._activeSession is not None:
                    try:
                        PySparkSession._activeSession.stop()
                    except Exception as e:
                        logger.warning(
                            "session.stop(): Local Spark Connect Server could not be stopped. "
                            f"Error: ${e}"
                        )
                del os.environ["SPARK_LOCAL_REMOTE"]
                del os.environ["SPARK_CONNECT_MODE_ENABLED"]
                if "SPARK_REMOTE" in os.environ:
                    del os.environ["SPARK_REMOTE"]

    def __enter__(self) -> "SparkSession":
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.

        .. versionadded:: 4.3.0

        Examples
        --------
        >>> with SparkSession.builder.remote("local[*]").getOrCreate() as session:
        ...     session.range(5).show()  # doctest: +SKIP
        +---+
        | id|
        +---+
        |  0|
        |  1|
        |  2|
        |  3|
        |  4|
        +---+
        """
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.

        .. versionadded:: 4.3.0

        Examples
        --------
        >>> with SparkSession.builder.remote("local[*]").getOrCreate() as session:
        ...     session.range(5).show()  # doctest: +SKIP
        +---+
        | id|
        +---+
        |  0|
        |  1|
        |  2|
        |  3|
        |  4|
        +---+
        """
        self.stop()

    @property
    def is_stopped(self) -> bool:
        """
        Returns if this session was stopped
        """
        return self.client.is_closed

    @property
    def conf(self) -> RuntimeConf:
        return RuntimeConf(self.client)

    conf.__doc__ = PySparkSession.conf.__doc__

    @property
    def streams(self) -> "StreamingQueryManager":
        if hasattr(self, "_sqm"):
            return self._sqm
        self._sqm: StreamingQueryManager = StreamingQueryManager(self)
        return self._sqm

    streams.__doc__ = PySparkSession.streams.__doc__

    def __getattr__(self, name: str) -> Any:
        if name in ["_jsc", "_jconf", "_jvm", "_jsparkSession", "sparkContext"]:
            raise PySparkAttributeError(
                errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": name}
            )
        return object.__getattribute__(self, name)

    @property
    def udf(self) -> "UDFRegistration":
        from pyspark.sql.connect.udf import UDFRegistration

        return UDFRegistration(self)

    udf.__doc__ = PySparkSession.udf.__doc__

    @property
    def udtf(self) -> "UDTFRegistration":
        from pyspark.sql.connect.udtf import UDTFRegistration

        return UDTFRegistration(self)

    udtf.__doc__ = PySparkSession.udtf.__doc__

    @property
    def dataSource(self) -> "DataSourceRegistration":
        from pyspark.sql.connect.datasource import DataSourceRegistration

        return DataSourceRegistration(self)

    dataSource.__doc__ = PySparkSession.dataSource.__doc__

    @functools.cached_property
    def version(self) -> str:
        result = self._client._analyze(method="spark_version").spark_version
        assert result is not None
        return result

    version.__doc__ = PySparkSession.version.__doc__

    @property
    def client(self) -> "SparkConnectClient":
        return self._client

    client.__doc__ = PySparkSession.client.__doc__

    def addArtifacts(
        self, *path: str, pyfile: bool = False, archive: bool = False, file: bool = False
    ) -> None:
        if sum([file, pyfile, archive]) > 1:
            raise PySparkValueError(
                errorClass="INVALID_MULTIPLE_ARGUMENT_CONDITIONS",
                messageParameters={
                    "arg_names": "'pyfile', 'archive' and/or 'file'",
                    "condition": "True together",
                },
            )
        self._client.add_artifacts(*path, pyfile=pyfile, archive=archive, file=file)

    addArtifacts.__doc__ = PySparkSession.addArtifacts.__doc__

    addArtifact = addArtifacts

    def _cache_local_relation(
        self,
        local_relation: LocalRelation,
        local_relation_size_limit: int,
        max_chunk_size_rows: int,
        max_chunk_size_bytes: int,
        max_batch_of_chunks_size_bytes: int,
    ) -> ChunkedCachedLocalRelation:
        """
        Cache the local relation at the server side if it has not been cached yet.

        This method serializes the input local relation into multiple data chunks and
        a schema chunk (if the schema is available) and uploads these chunks as artifacts
        to the server.

        The method collects a batch of chunks of size up to max_batch_of_chunks_size_bytes and
        uploads them together to the server.
        Uploading each chunk separately would require an additional RPC call for each chunk.
        Uploading all chunks together would require materializing all chunks in memory which
        may cause high memory usage on the client.
        Uploading batches of chunks is the middle-ground solution.

        Should only be called on a LocalRelation with a non-empty _table.
        """
        assert local_relation._table is not None, "table cannot be None"
        has_schema = local_relation._schema is not None

        hashes = []
        current_batch = []
        current_batch_size = 0
        total_size = 0
        if has_schema:
            schema_chunk = local_relation._serialize_schema()
            current_batch.append(schema_chunk)
            current_batch_size += len(schema_chunk)
            total_size += len(schema_chunk)

        data_chunks: Iterator[bytes] = local_relation._serialize_table_chunks(
            max_chunk_size_rows, min(max_chunk_size_bytes, max_batch_of_chunks_size_bytes)
        )

        for chunk in data_chunks:
            chunk_size = len(chunk)
            total_size += chunk_size

            # Check if total size exceeds the limit
            if total_size > local_relation_size_limit:
                raise AnalysisException(
                    errorClass="LOCAL_RELATION_SIZE_LIMIT_EXCEEDED",
                    messageParameters={
                        "actualSize": str(total_size),
                        "sizeLimit": str(local_relation_size_limit),
                    },
                )

            # Check if adding this chunk would exceed batch size
            if (
                len(current_batch) > 0
                and current_batch_size + chunk_size > max_batch_of_chunks_size_bytes
            ):
                hashes += self._client.cache_artifacts(current_batch)
                # start a new batch
                current_batch = []
                current_batch_size = 0

            current_batch.append(chunk)
            current_batch_size += chunk_size
        hashes += self._client.cache_artifacts(current_batch)

        if has_schema:
            schema_hash = hashes[0]
            data_hashes = hashes[1:]
        else:
            schema_hash = None
            data_hashes = hashes
        return ChunkedCachedLocalRelation(data_hashes, schema_hash)

    def copyFromLocalToFs(self, local_path: str, dest_path: str) -> None:
        if urllib.parse.urlparse(dest_path).scheme:
            raise PySparkValueError(
                errorClass="NO_SCHEMA_AND_DRIVER_DEFAULT_SCHEME",
                messageParameters={"arg_name": "dest_path"},
            )
        self._client.copy_from_local_to_fs(local_path, dest_path)

    copyFromLocalToFs.__doc__ = PySparkSession.copyFromLocalToFs.__doc__

    def _create_remote_dataframe(self, remote_id: str) -> "ParentDataFrame":
        """
        In internal API to reference a runtime DataFrame on the server side.
        This is used in ForeachBatch() runner, where the remote DataFrame refers to the
        output of a micro batch.
        """
        return DataFrame(CachedRemoteRelation(remote_id, spark_session=self), self)

    @staticmethod
    def _start_connect_server(master: str, opts: Dict[str, Any]) -> None:
        """
        Starts the Spark Connect server given the master (thread-unsafe).

        1. Temporarily remove all states for Spark Connect, for example, ``SPARK_REMOTE``
          environment variable.

        2. Starts a regular Spark session that automatically starts a Spark Connect server
           via ``spark.plugins`` feature.

        Returns the authentication token that should be used to connect to this session.
        """
        from pyspark import SparkContext, SparkConf

        session = PySparkSession._instantiatedSession
        if session is None or session._sc._jsc is None:
            init_opts = {}
            for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))):
                init_opts = json.loads(os.environ[f"PYSPARK_REMOTE_INIT_CONF_{i}"])
            init_opts.update(opts)
            opts = init_opts

            # Configurations to be overwritten
            overwrite_conf = opts
            overwrite_conf["spark.master"] = master
            if "spark.remote" in overwrite_conf:
                del overwrite_conf["spark.remote"]
            if "spark.api.mode" in overwrite_conf:
                del overwrite_conf["spark.api.mode"]

            # Check for a user provided authentication token, creating a new one if not,
            # and make sure it's set in the environment,
            if "SPARK_CONNECT_AUTHENTICATE_TOKEN" not in os.environ:
                os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = opts.get(
                    "spark.connect.authenticate.token", str(uuid.uuid4())
                )

            # Configurations to be set if unset.
            default_conf = {
                "spark.plugins": "org.apache.spark.sql.connect.SparkConnectPlugin",
                "spark.sql.artifact.isolation.enabled": "true",
                "spark.sql.artifact.isolation.alwaysApplyClassloader": "true",
            }

            if "SPARK_TESTING" in os.environ:
                # For testing, we use 0 to use an ephemeral port to allow parallel testing.
                # See also SPARK-42272.
                overwrite_conf["spark.connect.grpc.binding.port"] = "0"

            origin_remote = os.environ.get("SPARK_REMOTE", None)
            try:
                # So SparkSubmit thinks no remote is set in order to
                # start the regular PySpark session.
                if origin_remote is not None:
                    del os.environ["SPARK_REMOTE"]

                # The regular PySpark session is registered as an active session
                # so would not be garbage-collected.
                conf = SparkConf(loadDefaults=True)
                conf.setAll(list(overwrite_conf.items())).setAll(list(default_conf.items()))
                PySparkSession(SparkContext.getOrCreate(conf))

                # Lastly only keep runtime configurations because other configurations are
                # disallowed to set in the regular Spark Connect session.
                utl = SparkContext._jvm.PythonSQLUtils  # type: ignore[union-attr]
                runtime_conf_keys = [c._1() for c in utl.listRuntimeSQLConfigs()]
                new_opts = {k: opts[k] for k in opts if k in runtime_conf_keys}
                opts.clear()
                opts.update(new_opts)
            finally:
                if origin_remote is not None:
                    os.environ["SPARK_REMOTE"] = origin_remote
        else:
            raise PySparkRuntimeError(
                errorClass="SESSION_OR_CONTEXT_EXISTS",
                messageParameters={},
            )

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def _profiler_collector(self) -> ProfilerCollector:
        return self._client._profiler_collector

    @property
    def profile(self) -> Profile:
        return Profile(self._client._profiler_collector)

    profile.__doc__ = PySparkSession.profile.__doc__

    def __reduce__(self) -> Tuple:
        """
        This method is called when the object is pickled. It returns a tuple of the object's
        constructor function, arguments to it and the local state of the object.
        This function is supposed to only be used when the active spark session that is pickled
        is the same active spark session that is unpickled.
        """

        def creator(old_session_id: str) -> "SparkSession":
            # We cannot perform the checks for session matching here because accessing the
            # session ID property causes the serialization of a WeakRef and in turn breaks
            # the serialization.
            return SparkSession._getActiveSessionIfMatches(old_session_id)

        return creator, (self._session_id,)

    def _to_ddl(self, struct: StructType) -> str:
        ddl = self._client._analyze(method="json_to_ddl", json_string=struct.json()).ddl_string
        assert ddl is not None
        return ddl

    def _parse_ddl(self, ddl: str) -> DataType:
        dt = self._client._analyze(method="ddl_parse", ddl_string=ddl).parsed
        assert dt is not None
        return dt

    def cloneSession(self, new_session_id: Optional[str] = None) -> "SparkSession":
        """
        Create a clone of this Spark Connect session on the server side. The server-side session
        is cloned with all its current state (SQL configurations, temporary views, registered
        functions, catalog state) copied over to a new independent session. The returned cloned
        session is isolated from this session - any subsequent changes to either session's
        server-side state will not be reflected in the other.

        Parameters
        ----------
        new_session_id : str, optional
            Custom session ID to use for the cloned session (must be a valid UUID).
            If not provided, a new UUID will be generated.

        Returns
        -------
        SparkSession
            A new SparkSession instance with the cloned session.

        Notes
        -----
        This creates a new server-side session with the specified or generated session ID
        while preserving the current session's configuration and state.

        .. note::
            This is a developer API.
        """
        cloned_client = self._client.clone(new_session_id)
        # Create a new SparkSession with the cloned client directly
        new_session = object.__new__(SparkSession)
        new_session._client = cloned_client
        new_session._session_id = cloned_client._session_id
        new_session.release_session_on_close = True
        return new_session

    def newSession(self) -> "SparkSession":
        """
        Returns a new :class:`SparkSession` as a new session, that has separate SQLConf,
        registered temporary views and UDFs, but shared table cache.

        Unlike :meth:`cloneSession`, the returned session starts with empty state: no
        configuration, temporary views, registered functions, or catalog state are copied
        over from this session. This matches the Scala Connect ``newSession()`` semantics,
        which creates a completely fresh session against the same server. Note one
        difference from classic ``newSession()``: configurations set through
        ``SparkSession.builder.config(...)`` when this session was created are not
        reapplied to the new session; it starts from the server defaults.

        .. versionadded:: 4.3.0

        Returns
        -------
        :class:`SparkSession`
            A new SparkSession bound to a fresh, independent server-side session.
        """
        new_client = self._client.newSession()
        # Create a new SparkSession bound to the fresh, independent session directly.
        new_session = object.__new__(SparkSession)
        new_session._client = new_client
        new_session._session_id = new_client._session_id
        new_session.release_session_on_close = True
        return new_session


SparkSession.__doc__ = PySparkSession.__doc__


def _test() -> None:
    import os
    import sys
    import doctest
    from pyspark.sql import SparkSession as PySparkSession
    import pyspark.sql.connect.session

    globs = pyspark.sql.connect.session.__dict__.copy()
    globs["spark"] = (
        PySparkSession.builder.appName("sql.connect.session tests")
        .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
        .getOrCreate()
    )

    # Uses PySpark session to test builder.
    globs["SparkSession"] = PySparkSession
    # Spark Connect does not support to set master together.
    pyspark.sql.connect.session.SparkSession.__doc__ = None
    del pyspark.sql.connect.session.SparkSession.Builder.master.__doc__

    failure_count, test_count = doctest.testmod(
        pyspark.sql.connect.session,
        globs=globs,
        optionflags=doctest.ELLIPSIS
        | doctest.NORMALIZE_WHITESPACE
        | doctest.IGNORE_EXCEPTION_DETAIL,
    )

    globs["spark"].stop()

    if failure_count:
        sys.exit(-1)


if __name__ == "__main__":
    _test()