From 51f438bdfb3a68cc717e4812c7b5ee8251881ac4 Mon Sep 17 00:00:00 2001 From: Sagar Mittal Date: Wed, 10 Jun 2026 15:36:56 -0700 Subject: [PATCH 1/3] Make small optimizations to StatefulProcessorApiClient --- .../stateful_processor_api_client.py | 70 ++++++++++--------- 1 file changed, 38 insertions(+), 32 deletions(-) diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py b/python/pyspark/sql/streaming/stateful_processor_api_client.py index eab91e0c3f840..6ceed72936e5f 100644 --- a/python/pyspark/sql/streaming/stateful_processor_api_client.py +++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py @@ -14,11 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from datetime import datetime from enum import Enum import json import os import socket -from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast +from types import NoneType +from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast, TYPE_CHECKING from pyspark.serializers import write_int, read_int, UTF8Deserializer from pyspark.sql.pandas.serializers import ArrowStreamSerializer @@ -27,12 +29,43 @@ Row, ) from pyspark.sql.pandas.types import convert_pandas_using_numpy_type -from pyspark.serializers import CPickleSerializer +from pyspark.serializers import PickleSerializer from pyspark.errors import PySparkRuntimeError import uuid __all__ = ["StatefulProcessorApiClient", "StatefulProcessorHandleState"] +try: + import numpy as np + + has_numpy = True + + def _normalize_state_value(v: Any) -> Any: + # Fast path for common scalar values. + if isinstance(v, (bool, int, float, str, bytes, datetime, NoneType)): + return v + # Convert NumPy scalar values to Python primitive values. + if isinstance(v, np.generic): + return v.tolist() + # Named tuples (collections.namedtuple or typing.NamedTuple) and Row both + # require positional arguments and cannot be instantiated with a generator expression. + if isinstance(v, Row) or (isinstance(v, tuple) and hasattr(v, "_fields")): + return type(v)(*map(_normalize_state_value, v)) + # List / tuple: recursively normalize each element. + if isinstance(v, (list, tuple)): + return type(v)(map(_normalize_state_value, v)) + # Dict: normalize both keys and values. + if isinstance(v, dict): + return {_normalize_state_value(k): _normalize_state_value(val) for k, val in v.items()} + # Address a couple of pandas dtypes too. + if hasattr(v, "to_pytimedelta"): + return v.to_pytimedelta() + if hasattr(v, "to_pydatetime"): + return v.to_pydatetime() + return v +except ImportError: + has_numpy = False + class StatefulProcessorHandleState(Enum): PRE_INIT = 0 @@ -74,7 +107,7 @@ def __init__( else: self.handle_state = StatefulProcessorHandleState.CREATED self.utf8_deserializer = UTF8Deserializer() - self.pickleSer = CPickleSerializer() + self.pickleSer = PickleSerializer() self.serializer = ArrowStreamSerializer() # Dictionaries to store the mapping between iterator id and a tuple of data batch # and the index of the last row that was read. @@ -488,35 +521,8 @@ def _receive_str(self) -> str: return self.utf8_deserializer.loads(self.sockfile) def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes: - from pyspark.testing.utils import have_numpy - - if have_numpy: - import numpy as np - - def normalize_value(v: Any) -> Any: - # Convert NumPy types to Python primitive types. - if isinstance(v, np.generic): - return v.tolist() - # Named tuples (collections.namedtuple or typing.NamedTuple) and Row both - # require positional arguments and cannot be instantiated - # with a generator expression. - if isinstance(v, Row) or (isinstance(v, tuple) and hasattr(v, "_fields")): - return type(v)(*[normalize_value(e) for e in v]) - # List / tuple: recursively normalize each element - if isinstance(v, (list, tuple)): - return type(v)(normalize_value(e) for e in v) - # Dict: normalize both keys and values - if isinstance(v, dict): - return {normalize_value(k): normalize_value(val) for k, val in v.items()} - # Address a couple of pandas dtypes too. - elif hasattr(v, "to_pytimedelta"): - return v.to_pytimedelta() - elif hasattr(v, "to_pydatetime"): - return v.to_pydatetime() - else: - return v - - converted = tuple(normalize_value(v) for v in data) + if has_numpy: + converted = tuple(map(_normalize_state_value, data)) else: converted = data From 56e58866721388cf4578a6bac29f4fc5493ef9c3 Mon Sep 17 00:00:00 2001 From: Sagar Mittal Date: Fri, 26 Jun 2026 11:26:48 -0700 Subject: [PATCH 2/3] remove unused import --- python/pyspark/sql/streaming/stateful_processor_api_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py b/python/pyspark/sql/streaming/stateful_processor_api_client.py index 6ceed72936e5f..4a83002983165 100644 --- a/python/pyspark/sql/streaming/stateful_processor_api_client.py +++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py @@ -20,7 +20,7 @@ import os import socket from types import NoneType -from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast, TYPE_CHECKING +from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast from pyspark.serializers import write_int, read_int, UTF8Deserializer from pyspark.sql.pandas.serializers import ArrowStreamSerializer From a6ed0f328678cf163b339ee98d8c4b361bf719d9 Mon Sep 17 00:00:00 2001 From: Sagar Mittal Date: Fri, 26 Jun 2026 11:45:41 -0700 Subject: [PATCH 3/3] Perform an exact type check rather than instanceof Benchmarked a tuple vs a frozenset, and the tuple came out noticeably faster. --- .../pyspark/sql/streaming/stateful_processor_api_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py b/python/pyspark/sql/streaming/stateful_processor_api_client.py index 4a83002983165..166ac52fa6d87 100644 --- a/python/pyspark/sql/streaming/stateful_processor_api_client.py +++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py @@ -19,7 +19,6 @@ import json import os import socket -from types import NoneType from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast from pyspark.serializers import write_int, read_int, UTF8Deserializer @@ -39,10 +38,10 @@ import numpy as np has_numpy = True + SCALAR_TYPES = (bool, int, float, str, bytes, datetime, type(None)) def _normalize_state_value(v: Any) -> Any: - # Fast path for common scalar values. - if isinstance(v, (bool, int, float, str, bytes, datetime, NoneType)): + if type(v) in SCALAR_TYPES: # Fast path for common scalar values. return v # Convert NumPy scalar values to Python primitive values. if isinstance(v, np.generic):