This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 181acbd8db57 [SPARK-57760][SS] Make small optimizations to
StatefulProcessorApiClient
181acbd8db57 is described below
commit 181acbd8db5738ff37509357c1b4d8e82396db27
Author: Sagar Mittal <[email protected]>
AuthorDate: Thu Jul 9 13:35:59 2026 +0900
[SPARK-57760][SS] Make small optimizations to StatefulProcessorApiClient
### What changes were proposed in this pull request?
Make two small optimizations to StatefulProcessorApiClient:
1. Call PickleSerializer() instead of using the default CPickleSerializer
(which is CloudPickleSerializer). We don't need the latter since this path does
not deal with code objects.
2. Micro-optimize state value normalization: add a fast-path for
primitives, prefer `map` to generator comprehensions, and move the numpy
import and function definition to the top level so it is done once
#### Benchmarks
This is a [microbenchmark for
`_serialize_to_bytes`](https://gist.github.com/funrollloops/62bf5fae1654910b63a4539e1181db91):
##### Before
```
--- single-field tuples ---
python int → LongType p50= 4.11µs p95=
4.21µs p99= 6.49µs
python float → DoubleType p50= 4.10µs p95=
4.20µs p99= 4.96µs
np.float64 → DoubleType p50= 4.42µs p95=
4.57µs p99= 6.46µs
np.datetime64 → Timestamp p50= 7.32µs p95=
7.67µs p99= 11.60µs
pd.Timestamp → Timestamp p50= 7.53µs p95=
7.76µs p99= 12.34µs
--- wider tuples ---
10× python float → 10× DoubleType p50= 7.80µs p95=
7.94µs p99= 12.37µs
10× np.float64 → 10× DoubleType p50= 9.01µs p95=
9.23µs p99= 14.88µs
mixed (np.f64, np.i64, str, pd.Ts) p50= 9.93µs p95=
10.26µs p99= 17.66µs
```
##### After
```
--- single-field tuples ---
python int → LongType p50= 1.17µs p95=
1.19µs p99= 1.22µs
python float → DoubleType p50= 1.18µs p95=
1.19µs p99= 1.22µs
np.float64 → DoubleType p50= 1.92µs p95=
1.98µs p99= 2.01µs
np.datetime64 → Timestamp p50= 4.59µs p95=
4.71µs p99= 4.78µs
pd.Timestamp → Timestamp p50= 5.07µs p95=
5.17µs p99= 5.24µs
--- wider tuples ---
10× python float → 10× DoubleType p50= 2.19µs p95=
2.23µs p99= 2.26µs
10× np.float64 → 10× DoubleType p50= 7.72µs p95=
7.82µs p99= 7.89µs
mixed (np.f64, np.i64, str, pd.Ts) p50= 7.34µs p95=
7.46µs p99= 7.58µs
```
### Why are the changes needed?
Together these changes improve transform with state on a simple
rolling-window style benchmark by ~10%.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing unit tests.
### Was this patch authored or co-authored using generative AI tooling?
No, but Claude was consulted in the process of producing this PR.
Closes #56786 from funrollloops/tws-opt-1.
Lead-authored-by: Sagar Mittal <[email protected]>
Co-authored-by: Sagar Mittal <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
(cherry picked from commit 0af85eee81716bb2095c8ee2af39b77e7cca9a8f)
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../sql/streaming/stateful_processor_api_client.py | 67 ++++++++++++----------
1 file changed, 36 insertions(+), 31 deletions(-)
diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py
b/python/pyspark/sql/streaming/stateful_processor_api_client.py
index eab91e0c3f84..166ac52fa6d8 100644
--- a/python/pyspark/sql/streaming/stateful_processor_api_client.py
+++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py
@@ -14,6 +14,7 @@
# 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
@@ -27,12 +28,43 @@ from pyspark.sql.types import (
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
+ SCALAR_TYPES = (bool, int, float, str, bytes, datetime, type(None))
+
+ def _normalize_state_value(v: Any) -> Any:
+ 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):
+ 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 +106,7 @@ class StatefulProcessorApiClient:
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 +520,8 @@ class StatefulProcessorApiClient:
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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
