| # |
| # 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. |
| # |
| |
| """ |
| Serializers for PyArrow and pandas conversions. See `pyspark.serializers` for more details. |
| """ |
| |
| from itertools import groupby |
| from typing import IO, TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Tuple |
| |
| from pyspark.errors import PySparkRuntimeError, PySparkValueError |
| from pyspark.serializers import ( |
| Serializer, |
| read_int, |
| write_int, |
| UTF8Deserializer, |
| CPickleSerializer, |
| ) |
| from pyspark.sql import Row |
| from pyspark.sql.conversion import ( |
| ArrowBatchTransformer, |
| PandasToArrowConversion, |
| ) |
| from pyspark.sql.pandas.types import ( |
| from_arrow_schema, |
| to_arrow_type, |
| ) |
| from pyspark.sql.types import ( |
| DataType, |
| StringType, |
| StructType, |
| BinaryType, |
| StructField, |
| LongType, |
| IntegerType, |
| ) |
| |
| if TYPE_CHECKING: |
| import pandas as pd |
| import pyarrow as pa |
| |
| |
| def _normalize_packed(packed): |
| """ |
| Normalize UDF output to a uniform tuple-of-tuples form. |
| |
| Iterator UDFs yield a single (series, spark_type) tuple directly, |
| while batched UDFs return a tuple of tuples ((s1, t1), (s2, t2), ...). |
| This function normalizes both forms to a tuple of tuples. |
| """ |
| if len(packed) == 2 and isinstance(packed[1], DataType): |
| return (packed,) |
| return tuple(packed) |
| |
| |
| class SpecialLengths: |
| END_OF_DATA_SECTION = -1 |
| PYTHON_EXCEPTION_THROWN = -2 |
| TIMING_DATA = -3 |
| END_OF_STREAM = -4 |
| NULL = -5 |
| START_ARROW_STREAM = -6 |
| |
| |
| class ArrowCollectSerializer(Serializer): |
| """ |
| Deserialize a stream of batches followed by batch order information. Used in |
| PandasConversionMixin._collect_as_arrow() after invoking Dataset.collectAsArrowToPython() |
| in the JVM. |
| """ |
| |
| def __init__(self): |
| self.serializer = ArrowStreamSerializer() |
| |
| def dump_stream(self, iterator, stream): |
| return self.serializer.dump_stream(iterator, stream) |
| |
| def load_stream(self, stream): |
| """ |
| Load a stream of un-ordered Arrow RecordBatches, where the last iteration yields |
| a list of indices that can be used to put the RecordBatches in the correct order. |
| """ |
| # load the batches |
| for batch in self.serializer.load_stream(stream): |
| yield batch |
| |
| # load the batch order indices or propagate any error that occurred in the JVM |
| num = read_int(stream) |
| if num == -1: |
| error_msg = UTF8Deserializer().loads(stream) |
| raise PySparkRuntimeError( |
| errorClass="ERROR_OCCURRED_WHILE_CALLING", |
| messageParameters={ |
| "func_name": "ArrowCollectSerializer.load_stream", |
| "error_msg": error_msg, |
| }, |
| ) |
| batch_order = [] |
| for i in range(num): |
| index = read_int(stream) |
| batch_order.append(index) |
| yield batch_order |
| |
| def __repr__(self): |
| return "ArrowCollectSerializer(%s)" % self.serializer |
| |
| |
| class ArrowStreamSerializer(Serializer): |
| """ |
| Serializes Arrow record batches as a plain stream. |
| |
| Parameters |
| ---------- |
| write_start_stream : bool |
| If True, writes the START_ARROW_STREAM marker before the first |
| output batch. Default False. |
| """ |
| |
| def __init__(self, write_start_stream: bool = False) -> None: |
| super().__init__() |
| self._write_start_stream: bool = write_start_stream |
| |
| def dump_stream(self, iterator: Iterable["pa.RecordBatch"], stream: IO[bytes]) -> None: |
| """Optionally prepend START_ARROW_STREAM, then write batches.""" |
| iterator = iter(iterator) |
| if self._write_start_stream: |
| iterator = self._write_stream_start(iterator, stream) |
| import pyarrow as pa |
| |
| writer = None |
| try: |
| for batch in iterator: |
| if writer is None: |
| writer = pa.RecordBatchStreamWriter(stream, batch.schema) |
| writer.write_batch(batch) |
| finally: |
| if writer is not None: |
| writer.close() |
| |
| def load_stream(self, stream: IO[bytes]) -> Iterator["pa.RecordBatch"]: |
| """Load batches from a plain Arrow stream.""" |
| import pyarrow as pa |
| |
| reader = pa.ipc.open_stream(stream) |
| for batch in reader: |
| yield batch |
| |
| def _write_stream_start( |
| self, batch_iterator: Iterator["pa.RecordBatch"], stream: IO[bytes] |
| ) -> Iterator["pa.RecordBatch"]: |
| """Write START_ARROW_STREAM before the first batch, then pass batches through.""" |
| import itertools |
| |
| first = next(batch_iterator, None) |
| if first is None: |
| return |
| |
| # Signal the JVM after the first batch succeeds, so errors during |
| # batch creation can be reported before the Arrow stream starts. |
| write_int(SpecialLengths.START_ARROW_STREAM, stream) |
| yield from itertools.chain([first], batch_iterator) |
| |
| def __repr__(self) -> str: |
| return "ArrowStreamSerializer(write_start_stream=%s)" % self._write_start_stream |
| |
| |
| class ArrowStreamGroupSerializer(ArrowStreamSerializer): |
| """ |
| Extends :class:`ArrowStreamSerializer` with group-count protocol for loading |
| grouped Arrow record batches (1 dataframe per group). |
| """ |
| |
| def load_stream(self, stream: IO[bytes]) -> Iterator[Iterator["pa.RecordBatch"]]: |
| """Yield one iterator of record batches per group from the stream.""" |
| while dataframes_in_group := read_int(stream): |
| if dataframes_in_group == 1: |
| yield ArrowStreamSerializer.load_stream(self, stream) |
| elif dataframes_in_group > 0: |
| raise PySparkValueError( |
| errorClass="INVALID_NUMBER_OF_DATAFRAMES_IN_GROUP", |
| messageParameters={"dataframes_in_group": str(dataframes_in_group)}, |
| ) |
| |
| |
| class ArrowStreamCoGroupSerializer(ArrowStreamSerializer): |
| """ |
| Extends :class:`ArrowStreamSerializer` with group-count protocol for loading |
| cogrouped Arrow record batches (2 dataframes per group). |
| """ |
| |
| def load_stream( |
| self, stream: IO[bytes] |
| ) -> Iterator[Tuple[List["pa.RecordBatch"], List["pa.RecordBatch"]]]: |
| """Yield pairs of (left_batches, right_batches) from the stream.""" |
| while dataframes_in_group := read_int(stream): |
| if dataframes_in_group == 2: |
| # Must eagerly load each dataframe to maintain correct stream position |
| yield ( |
| list(ArrowStreamSerializer.load_stream(self, stream)), |
| list(ArrowStreamSerializer.load_stream(self, stream)), |
| ) |
| elif dataframes_in_group > 0: |
| raise PySparkValueError( |
| errorClass="INVALID_NUMBER_OF_DATAFRAMES_IN_GROUP", |
| messageParameters={"dataframes_in_group": str(dataframes_in_group)}, |
| ) |
| |
| |
| class ArrowStreamUDFSerializer(ArrowStreamSerializer): |
| """ |
| Same as :class:`ArrowStreamSerializer` but it flattens the struct to Arrow record batch |
| for applying each function with the raw record arrow batch. See also `DataFrame.mapInArrow`. |
| """ |
| |
| def load_stream(self, stream): |
| """ |
| Flatten the struct into Arrow's record batches. |
| """ |
| batches = super().load_stream(stream) |
| flattened = map(ArrowBatchTransformer.flatten_struct, batches) |
| return map(lambda b: [b], flattened) |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Override because Pandas UDFs require a START_ARROW_STREAM before the Arrow stream is sent. |
| """ |
| batches = self._write_stream_start( |
| (ArrowBatchTransformer.wrap_struct(x[0]) for x in iterator), stream |
| ) |
| return super().dump_stream(batches, stream) |
| |
| |
| class ArrowStreamUDTFSerializer(ArrowStreamUDFSerializer): |
| """ |
| Same as :class:`ArrowStreamUDFSerializer` but it does not flatten when loading batches. |
| """ |
| |
| def load_stream(self, stream): |
| return ArrowStreamSerializer.load_stream(self, stream) |
| |
| |
| class ArrowStreamArrowUDTFSerializer(ArrowStreamUDTFSerializer): |
| """ |
| Serializer for PyArrow-native UDTFs that work directly with PyArrow RecordBatches and Arrays. |
| """ |
| |
| def __init__(self, *, table_arg_offsets=None): |
| super().__init__() |
| self.table_arg_offsets = table_arg_offsets if table_arg_offsets else [] |
| |
| def load_stream(self, stream): |
| """ |
| Flatten the struct into Arrow's record batches. |
| """ |
| for batch in super().load_stream(stream): |
| # For each column: flatten struct columns at table_arg_offsets into RecordBatch, |
| # keep other columns as Array |
| yield [ |
| ( |
| ArrowBatchTransformer.flatten_struct(batch, column_index=i) |
| if i in self.table_arg_offsets |
| else batch.column(i) |
| ) |
| for i in range(batch.num_columns) |
| ] |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Override to handle type coercion for ArrowUDTF outputs. |
| ArrowUDTF returns iterator of (pa.RecordBatch, arrow_return_type) tuples. |
| """ |
| import pyarrow as pa |
| |
| def apply_type_coercion(): |
| for batch, arrow_return_type in iterator: |
| assert isinstance(arrow_return_type, pa.StructType), ( |
| f"Expected pa.StructType, got {type(arrow_return_type)}" |
| ) |
| coerced_batch = ArrowBatchTransformer.enforce_schema( |
| batch, pa.schema(arrow_return_type), safecheck=True |
| ) |
| yield coerced_batch, arrow_return_type |
| |
| return super().dump_stream(apply_type_coercion(), stream) |
| |
| |
| class ArrowStreamGroupUDFSerializer(ArrowStreamUDFSerializer): |
| """ |
| Serializer for grouped Arrow UDFs. |
| |
| Deserializes: |
| ``Iterator[Iterator[pa.RecordBatch]]`` - one inner iterator per group. |
| Each batch contains a single struct column. |
| |
| Serializes: |
| ``Iterator[Tuple[Iterator[pa.RecordBatch], pa.DataType]]`` |
| Each tuple contains iterator of flattened batches and their Arrow type. |
| |
| Used by: |
| - SQL_GROUPED_MAP_ARROW_UDF |
| - SQL_GROUPED_MAP_ARROW_ITER_UDF |
| |
| Parameters |
| ---------- |
| assign_cols_by_name : bool |
| If True, reorder serialized columns by schema name. |
| """ |
| |
| def __init__(self, *, assign_cols_by_name): |
| super().__init__() |
| self._assign_cols_by_name = assign_cols_by_name |
| |
| def load_stream(self, stream): |
| """ |
| Load grouped Arrow record batches from stream. |
| """ |
| for batches in ArrowStreamGroupSerializer.load_stream(self, stream): |
| yield batches |
| # Make sure the batches are fully iterated before getting the next group |
| for _ in batches: |
| pass |
| |
| def dump_stream(self, iterator, stream): |
| import pyarrow as pa |
| |
| # flatten inner list [([pa.RecordBatch], arrow_type)] into [(pa.RecordBatch, arrow_type)] |
| # so strip off inner iterator induced by ArrowStreamUDFSerializer.load_stream |
| batch_iter = ( |
| (batch, arrow_type) |
| for batches, arrow_type in iterator # tuple constructed in wrap_grouped_map_arrow_udf |
| for batch in batches |
| ) |
| |
| if self._assign_cols_by_name: |
| batch_iter = ( |
| ( |
| pa.RecordBatch.from_arrays( |
| [batch.column(field.name) for field in arrow_type], |
| names=[field.name for field in arrow_type], |
| ), |
| arrow_type, |
| ) |
| for batch, arrow_type in batch_iter |
| ) |
| |
| super().dump_stream(batch_iter, stream) |
| |
| |
| class ArrowStreamPandasSerializer(ArrowStreamSerializer): |
| """ |
| Serializes pandas.Series as Arrow data with Arrow streaming format. |
| |
| Parameters |
| ---------- |
| timezone : str |
| A timezone to respect when handling timestamp values |
| safecheck : bool |
| If True, conversion from Arrow to Pandas checks for overflow/truncation |
| int_to_decimal_coercion_enabled : bool |
| If True, applies additional coercions in Python before converting to Arrow. |
| This has performance penalties. |
| prefers_large_types : bool |
| If True, prefer large Arrow types (e.g., large_string instead of string). |
| struct_in_pandas : str, optional |
| How to represent struct in pandas ("dict", "row", etc.). Default is "dict". |
| ndarray_as_list : bool, optional |
| Whether to convert ndarray as list. Default is False. |
| prefer_int_ext_dtype : bool, optional |
| Whether to convert integers to Pandas ExtensionDType. Default is False. |
| df_for_struct : bool, optional |
| If True, convert struct columns to DataFrame instead of Series. Default is False. |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| int_to_decimal_coercion_enabled: bool = False, |
| prefers_large_types: bool = False, |
| struct_in_pandas: str = "dict", |
| ndarray_as_list: bool = False, |
| prefer_int_ext_dtype: bool = False, |
| df_for_struct: bool = False, |
| input_type: Optional["StructType"] = None, |
| arrow_cast: bool = False, |
| ): |
| super().__init__() |
| self._timezone = timezone |
| self._safecheck = safecheck |
| self._int_to_decimal_coercion_enabled = int_to_decimal_coercion_enabled |
| self._prefers_large_types = prefers_large_types |
| self._struct_in_pandas = struct_in_pandas |
| self._ndarray_as_list = ndarray_as_list |
| self._prefer_int_ext_dtype = prefer_int_ext_dtype |
| self._df_for_struct = df_for_struct |
| if input_type is not None: |
| assert isinstance(input_type, StructType) |
| self._input_type = input_type |
| self._arrow_cast = arrow_cast |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Make ArrowRecordBatches from Pandas Series and serialize. |
| Each element in iterator is: |
| - For batched UDFs: tuple of (series, spark_type) tuples: ((s1, t1), (s2, t2), ...) |
| - For iterator UDFs: single (series, spark_type) tuple directly |
| """ |
| |
| def create_batch( |
| series_tuples: Tuple[Tuple["pd.Series", DataType], ...], |
| ) -> "pa.RecordBatch": |
| series_data = [s for s, _ in series_tuples] |
| types = [t for _, t in series_tuples] |
| schema = StructType([StructField(f"_{i}", t) for i, t in enumerate(types)]) |
| return PandasToArrowConversion.convert( |
| series_data, |
| schema, |
| timezone=self._timezone, |
| safecheck=self._safecheck, |
| prefers_large_types=self._prefers_large_types, |
| int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, |
| ) |
| |
| super().dump_stream( |
| (create_batch(_normalize_packed(packed)) for packed in iterator), stream |
| ) |
| |
| def load_stream(self, stream): |
| """ |
| Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series. |
| """ |
| yield from map( |
| lambda batch: ArrowBatchTransformer.to_pandas( |
| batch, |
| timezone=self._timezone, |
| schema=self._input_type, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ), |
| super().load_stream(stream), |
| ) |
| |
| def __repr__(self): |
| return "ArrowStreamPandasSerializer" |
| |
| |
| class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer): |
| """ |
| Serializer used by Python worker to evaluate Pandas UDFs |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| df_for_struct: bool = False, |
| struct_in_pandas: str = "dict", |
| ndarray_as_list: bool = False, |
| prefer_int_ext_dtype: bool = False, |
| arrow_cast: bool = False, |
| input_type: Optional[StructType] = None, |
| int_to_decimal_coercion_enabled: bool = False, |
| prefers_large_types: bool = False, |
| ignore_unexpected_complex_type_values: bool = False, |
| is_legacy: bool = False, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| prefers_large_types=prefers_large_types, |
| struct_in_pandas=struct_in_pandas, |
| ndarray_as_list=ndarray_as_list, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| df_for_struct=df_for_struct, |
| input_type=input_type, |
| arrow_cast=arrow_cast, |
| ) |
| self._assign_cols_by_name = assign_cols_by_name |
| self._ignore_unexpected_complex_type_values = ignore_unexpected_complex_type_values |
| self._is_legacy = is_legacy |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Override because Pandas UDFs require a START_ARROW_STREAM before the Arrow stream is sent. |
| This should be sent after creating the first record batch so in case of an error, it can |
| be sent back to the JVM before the Arrow stream starts. |
| |
| Each element in iterator is: |
| - For batched UDFs: tuple of (series, spark_type) tuples: ((s1, t1), (s2, t2), ...) |
| - For iterator UDFs: single (series, spark_type) tuple directly |
| """ |
| import pandas as pd |
| |
| def create_batch( |
| series_tuples: Tuple[Tuple["pd.Series", DataType], ...], |
| ) -> "pa.RecordBatch": |
| # When struct_in_pandas="dict", UDF must return DataFrame for struct types |
| if self._struct_in_pandas == "dict": |
| for s, spark_type in series_tuples: |
| if isinstance(spark_type, StructType) and not isinstance(s, pd.DataFrame): |
| raise PySparkValueError( |
| "Invalid return type. Please make sure that the UDF returns a " |
| "pandas.DataFrame when the specified return type is StructType." |
| ) |
| |
| series_data = [s for s, _ in series_tuples] |
| types = [t for _, t in series_tuples] |
| schema = StructType([StructField(f"_{i}", t) for i, t in enumerate(types)]) |
| return PandasToArrowConversion.convert( |
| series_data, |
| schema, |
| timezone=self._timezone, |
| safecheck=self._safecheck, |
| arrow_cast=self._arrow_cast, |
| prefers_large_types=self._prefers_large_types, |
| assign_cols_by_name=self._assign_cols_by_name, |
| int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, |
| ignore_unexpected_complex_type_values=self._ignore_unexpected_complex_type_values, |
| is_legacy=self._is_legacy, |
| ) |
| |
| batches = self._write_stream_start( |
| (create_batch(_normalize_packed(packed)) for packed in iterator), |
| stream, |
| ) |
| return ArrowStreamSerializer.dump_stream(self, batches, stream) |
| |
| def __repr__(self): |
| return "ArrowStreamPandasUDFSerializer" |
| |
| |
| class ArrowStreamPandasUDTFSerializer(ArrowStreamPandasUDFSerializer): |
| """ |
| Serializer used by Python worker to evaluate Arrow-optimized Python UDTFs. |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| input_type, |
| prefer_int_ext_dtype, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| # The output pandas DataFrame's columns are unnamed. |
| assign_cols_by_name=False, |
| # Set to 'False' to avoid converting struct type inputs into a pandas DataFrame. |
| df_for_struct=False, |
| # Defines how struct type inputs are converted. If set to "row", struct type inputs |
| # are converted into Rows. Without this setting, a struct type input would be treated |
| # as a dictionary. For example, for named_struct('name', 'Alice', 'age', 1), |
| # if struct_in_pandas="dict", it becomes {"name": "Alice", "age": 1} |
| # if struct_in_pandas="row", it becomes Row(name="Alice", age=1) |
| struct_in_pandas="row", |
| # When dealing with array type inputs, Arrow converts them into numpy.ndarrays. |
| # To ensure consistency across regular and arrow-optimized UDTFs, we further |
| # convert these numpy.ndarrays into Python lists. |
| ndarray_as_list=True, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| # Enables explicit casting for mismatched return types of Arrow Python UDTFs. |
| arrow_cast=True, |
| input_type=input_type, |
| # Enable additional coercions for UDTF serialization |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| # UDTF-specific: ignore unexpected complex type values in converter |
| ignore_unexpected_complex_type_values=True, |
| # Legacy UDTF pandas conversion: enables broader Arrow exception |
| # handling to allow more implicit type coercions |
| is_legacy=True, |
| ) |
| |
| def __repr__(self): |
| return "ArrowStreamPandasUDTFSerializer" |
| |
| |
| # Serializer for SQL_GROUPED_AGG_PANDAS_UDF, SQL_WINDOW_AGG_PANDAS_UDF, |
| # and SQL_GROUPED_AGG_PANDAS_ITER_UDF |
| class ArrowStreamAggPandasUDFSerializer(ArrowStreamPandasUDFSerializer): |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| prefer_int_ext_dtype, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| assign_cols_by_name=assign_cols_by_name, |
| df_for_struct=False, |
| struct_in_pandas="dict", |
| ndarray_as_list=False, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| arrow_cast=True, |
| input_type=None, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| ) |
| |
| def load_stream(self, stream): |
| """ |
| Yield an iterator that produces one tuple of pandas.Series per batch. |
| Each group yields Iterator[Tuple[pd.Series, ...]], allowing UDF to |
| process batches one by one without consuming all batches upfront. |
| """ |
| for batches in ArrowStreamGroupSerializer.load_stream(self, stream): |
| # Lazily read and convert Arrow batches to pandas Series one at a time |
| # from the stream. This avoids loading all batches into memory for the group |
| series_iter = map( |
| lambda batch: tuple( |
| ArrowBatchTransformer.to_pandas( |
| batch, |
| timezone=self._timezone, |
| schema=self._input_type, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ) |
| ), |
| batches, |
| ) |
| yield series_iter |
| # Make sure the batches are fully iterated before getting the next group |
| for _ in series_iter: |
| pass |
| |
| def __repr__(self): |
| return "ArrowStreamAggPandasUDFSerializer" |
| |
| |
| # Serializer for SQL_GROUPED_MAP_PANDAS_UDF, SQL_GROUPED_MAP_PANDAS_ITER_UDF |
| class GroupPandasUDFSerializer(ArrowStreamPandasUDFSerializer): |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| prefer_int_ext_dtype, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| assign_cols_by_name=assign_cols_by_name, |
| df_for_struct=False, |
| struct_in_pandas="dict", |
| ndarray_as_list=False, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| arrow_cast=True, |
| input_type=None, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| ) |
| |
| def load_stream(self, stream): |
| """ |
| Deserialize Grouped ArrowRecordBatches and yield raw Iterator[pa.RecordBatch]. |
| Each outer iterator element represents a group. |
| """ |
| for batches in ArrowStreamGroupSerializer.load_stream(self, stream): |
| yield batches |
| # Make sure the batches are fully iterated before getting the next group |
| for _ in batches: |
| pass |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Flatten the grouped iterator structure. |
| """ |
| # Flatten: Iterator[Iterator[[(df, spark_type)]]] -> Iterator[[(df, spark_type)]] |
| flattened_iter = (batch for generator in iterator for batch in generator) |
| super().dump_stream(flattened_iter, stream) |
| |
| def __repr__(self): |
| return "GroupPandasUDFSerializer" |
| |
| |
| class CogroupArrowUDFSerializer(ArrowStreamGroupUDFSerializer): |
| """ |
| Serializes pyarrow.RecordBatch data with Arrow streaming format. |
| |
| Loads Arrow record batches as `[([pa.RecordBatch], [pa.RecordBatch])]` (one tuple per group) |
| and serializes `[([pa.RecordBatch], arrow_type)]`. |
| |
| Parameters |
| ---------- |
| assign_cols_by_name : bool |
| If True, then DataFrames will get columns by name |
| """ |
| |
| def load_stream(self, stream): |
| """ |
| Deserialize Cogrouped ArrowRecordBatches and yield as two `pyarrow.RecordBatch`es. |
| """ |
| yield from ArrowStreamCoGroupSerializer.load_stream(self, stream) |
| |
| |
| class CogroupPandasUDFSerializer(ArrowStreamPandasUDFSerializer): |
| def load_stream(self, stream): |
| """ |
| Deserialize Cogrouped ArrowRecordBatches to a tuple of Arrow tables and yield as two |
| lists of pandas.Series. |
| """ |
| import pyarrow as pa |
| |
| for left_batches, right_batches in ArrowStreamCoGroupSerializer.load_stream(self, stream): |
| left_table = pa.Table.from_batches(left_batches) |
| right_table = pa.Table.from_batches(right_batches) |
| yield ( |
| ArrowBatchTransformer.to_pandas( |
| left_table, |
| timezone=self._timezone, |
| schema=from_arrow_schema(left_table.schema), |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ), |
| ArrowBatchTransformer.to_pandas( |
| right_table, |
| timezone=self._timezone, |
| schema=from_arrow_schema(right_table.schema), |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ), |
| ) |
| |
| |
| class ApplyInPandasWithStateSerializer(ArrowStreamPandasUDFSerializer): |
| """ |
| Serializer used by Python worker to evaluate UDF for applyInPandasWithState. |
| |
| Parameters |
| ---------- |
| timezone : str |
| A timezone to respect when handling timestamp values |
| safecheck : bool |
| If True, conversion from Arrow to Pandas checks for overflow/truncation |
| assign_cols_by_name : bool |
| If True, then Pandas DataFrames will get columns by name |
| state_object_schema : StructType |
| The type of state object represented as Spark SQL type |
| arrow_max_records_per_batch : int |
| Limit of the number of records that can be written to a single ArrowRecordBatch in memory. |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| prefer_int_ext_dtype, |
| state_object_schema, |
| arrow_max_records_per_batch, |
| prefers_large_var_types, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| assign_cols_by_name=assign_cols_by_name, |
| df_for_struct=False, |
| struct_in_pandas="dict", |
| ndarray_as_list=False, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| arrow_cast=True, |
| input_type=None, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| prefers_large_types=prefers_large_var_types, |
| ) |
| self.pickleSer = CPickleSerializer() |
| self.utf8_deserializer = UTF8Deserializer() |
| self.state_object_schema = state_object_schema |
| |
| self.result_count_df_type = StructType( |
| [ |
| StructField("dataCount", IntegerType()), |
| StructField("stateCount", IntegerType()), |
| ] |
| ) |
| |
| self.result_count_pdf_arrow_type = to_arrow_type( |
| self.result_count_df_type, timezone="UTC", prefers_large_types=prefers_large_var_types |
| ) |
| |
| self.result_state_df_type = StructType( |
| [ |
| StructField("properties", StringType()), |
| StructField("keyRowAsUnsafe", BinaryType()), |
| StructField("object", BinaryType()), |
| StructField("oldTimeoutTimestamp", LongType()), |
| ] |
| ) |
| |
| self.result_state_pdf_arrow_type = to_arrow_type( |
| self.result_state_df_type, timezone="UTC", prefers_large_types=prefers_large_var_types |
| ) |
| self.arrow_max_records_per_batch = ( |
| arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 |
| ) |
| |
| def load_stream(self, stream): |
| """ |
| Read ArrowRecordBatches from stream, deserialize them to populate a list of pair |
| (data chunk, state), and convert the data into a list of pandas.Series. |
| |
| Please refer the doc of inner function `gen_data_and_state` for more details how |
| this function works in overall. |
| |
| In addition, this function further groups the return of `gen_data_and_state` by the state |
| instance (same semantic as grouping by grouping key) and produces an iterator of data |
| chunks for each group, so that the caller can lazily materialize the data chunk. |
| """ |
| |
| import pyarrow as pa |
| import json |
| from itertools import groupby |
| from pyspark.sql.streaming.state import GroupState |
| |
| def construct_state(state_info_col): |
| """ |
| Construct state instance from the value of state information column. |
| """ |
| |
| state_info_col_properties = state_info_col["properties"] |
| state_info_col_key_row = state_info_col["keyRowAsUnsafe"] |
| state_info_col_object = state_info_col["object"] |
| |
| state_properties = json.loads(state_info_col_properties) |
| if state_info_col_object: |
| state_object = self.pickleSer.loads(state_info_col_object) |
| else: |
| state_object = None |
| state_properties["optionalValue"] = state_object |
| |
| return GroupState( |
| keyAsUnsafe=state_info_col_key_row, |
| valueSchema=self.state_object_schema, |
| **state_properties, |
| ) |
| |
| def gen_data_and_state(batches): |
| """ |
| Deserialize ArrowRecordBatches and return a generator of |
| `(a list of pandas.Series, state)`. |
| |
| The logic on deserialization is following: |
| |
| 1. Read the entire data part from Arrow RecordBatch. |
| 2. Read the entire state information part from Arrow RecordBatch. |
| 3. Loop through each state information: |
| 3.A. Extract the data out from entire data via the information of data range. |
| 3.B. Construct a new state instance if the state information is the first occurrence |
| for the current grouping key. |
| 3.C. Leverage the existing state instance if it is already available for the current |
| grouping key. (Meaning it's not the first occurrence.) |
| 3.D. Remove the cache of state instance if the state information denotes the data is |
| the last chunk for current grouping key. |
| |
| This deserialization logic assumes that Arrow RecordBatches contain the data with the |
| ordering that data chunks for same grouping key will appear sequentially. |
| |
| This function must avoid materializing multiple Arrow RecordBatches into memory at the |
| same time. And data chunks from the same grouping key should appear sequentially, to |
| further group them based on state instance (same state instance will be produced for |
| same grouping key). |
| """ |
| |
| state_for_current_group = None |
| |
| for batch in batches: |
| batch_schema = batch.schema |
| data_schema = pa.schema([batch_schema[i] for i in range(0, len(batch_schema) - 1)]) |
| state_schema = pa.schema( |
| [ |
| batch_schema[-1], |
| ] |
| ) |
| |
| batch_columns = batch.columns |
| data_columns = batch_columns[0:-1] |
| state_column = batch_columns[-1] |
| |
| data_batch = pa.RecordBatch.from_arrays(data_columns, schema=data_schema) |
| state_batch = pa.RecordBatch.from_arrays( |
| [ |
| state_column, |
| ], |
| schema=state_schema, |
| ) |
| |
| state_pandas = ArrowBatchTransformer.to_pandas( |
| state_batch, |
| timezone=self._timezone, |
| schema=None, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| )[0] |
| |
| for state_idx in range(0, len(state_pandas)): |
| state_info_col = state_pandas.iloc[state_idx] |
| |
| if not state_info_col: |
| # no more data with grouping key + state |
| break |
| |
| data_start_offset = state_info_col["startOffset"] |
| num_data_rows = state_info_col["numRows"] |
| is_last_chunk = state_info_col["isLastChunk"] |
| |
| if state_for_current_group: |
| # use the state, we already have state for same group and there should be |
| # some data in same group being processed earlier |
| state = state_for_current_group |
| else: |
| # there is no state being stored for same group, construct one |
| state = construct_state(state_info_col) |
| |
| if is_last_chunk: |
| # discard the state being cached for same group |
| state_for_current_group = None |
| elif not state_for_current_group: |
| # there's no cached state but expected to have additional data in same group |
| # cache the current state |
| state_for_current_group = state |
| |
| data_batch_for_group = data_batch.slice(data_start_offset, num_data_rows) |
| data_pandas = ArrowBatchTransformer.to_pandas( |
| data_batch_for_group, |
| timezone=self._timezone, |
| schema=None, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ) |
| |
| # state info |
| yield ( |
| data_pandas, |
| state, |
| ) |
| |
| _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) |
| |
| data_state_generator = gen_data_and_state(_batches) |
| |
| # state will be same object for same grouping key |
| for _state, _data in groupby(data_state_generator, key=lambda x: x[1]): |
| yield ( |
| _data, |
| _state, |
| ) |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Read through an iterator of (iterator of pandas DataFrame, state), serialize them to Arrow |
| RecordBatches, and write batches to stream. |
| """ |
| |
| import pandas as pd |
| |
| def construct_state_pdf(state): |
| """ |
| Construct a pandas DataFrame from the state instance. |
| """ |
| |
| state_properties = state.json().encode("utf-8") |
| state_key_row_as_binary = state._keyAsUnsafe |
| if state.exists: |
| state_object = self.pickleSer.dumps(state._value_schema.toInternal(state._value)) |
| else: |
| state_object = None |
| state_old_timeout_timestamp = state.oldTimeoutTimestamp |
| |
| state_dict = { |
| "properties": [ |
| state_properties, |
| ], |
| "keyRowAsUnsafe": [ |
| state_key_row_as_binary, |
| ], |
| "object": [ |
| state_object, |
| ], |
| "oldTimeoutTimestamp": [ |
| state_old_timeout_timestamp, |
| ], |
| } |
| |
| return pd.DataFrame.from_dict(state_dict) |
| |
| def construct_record_batch(pdfs, pdf_data_cnt, pdf_schema, state_pdfs, state_data_cnt): |
| """ |
| Construct a new Arrow RecordBatch based on output pandas DataFrames and states. Each |
| one matches to the single struct field for Arrow schema. We also need an extra one to |
| indicate array length for data and state, so the return value of Arrow RecordBatch will |
| have schema with three fields, in `count`, `data`, `state` order. |
| (Readers are expected to access the field via position rather than the name. We do |
| not guarantee the name of the field.) |
| |
| Note that Arrow RecordBatch requires all columns to have all same number of rows, |
| hence this function inserts empty data for count/state/data with less elements to |
| compensate. |
| """ |
| |
| max_data_cnt = max(1, max(pdf_data_cnt, state_data_cnt)) |
| |
| # We only use the first row in the count column, and fill other rows to be the same |
| # value, hoping it is more friendly for compression, in case it is needed. |
| count_dict = { |
| "dataCount": [pdf_data_cnt] * max_data_cnt, |
| "stateCount": [state_data_cnt] * max_data_cnt, |
| } |
| count_pdf = pd.DataFrame.from_dict(count_dict) |
| |
| empty_row_cnt_in_data = max_data_cnt - pdf_data_cnt |
| empty_row_cnt_in_state = max_data_cnt - state_data_cnt |
| |
| empty_rows_pdf = pd.DataFrame( |
| dict.fromkeys(pdf_schema.names), |
| index=[x for x in range(0, empty_row_cnt_in_data)], |
| ) |
| empty_rows_state = pd.DataFrame( |
| columns=["properties", "keyRowAsUnsafe", "object", "oldTimeoutTimestamp"], |
| index=[x for x in range(0, empty_row_cnt_in_state)], |
| ) |
| |
| pdfs.append(empty_rows_pdf) |
| state_pdfs.append(empty_rows_state) |
| |
| merged_pdf = pd.concat(pdfs, ignore_index=True) |
| merged_state_pdf = pd.concat(state_pdfs, ignore_index=True) |
| |
| # Create batch from list of DataFrames, each wrapped as a StructArray. |
| # Schema fields map to: _0=count, _1=output data, _2=state data |
| # (types defined in __init__: result_count_df_type, pdf_schema, result_state_df_type) |
| data = [count_pdf, merged_pdf, merged_state_pdf] |
| schema = StructType( |
| [ |
| StructField("_0", self.result_count_df_type), |
| StructField("_1", pdf_schema), |
| StructField("_2", self.result_state_df_type), |
| ] |
| ) |
| return PandasToArrowConversion.convert( |
| data, |
| schema, |
| timezone=self._timezone, |
| safecheck=self._safecheck, |
| arrow_cast=self._arrow_cast, |
| assign_cols_by_name=self._assign_cols_by_name, |
| int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, |
| ) |
| |
| def serialize_batches(): |
| """ |
| Read through an iterator of (iterator of pandas DataFrame, state), and serialize them |
| to Arrow RecordBatches. |
| |
| This function does batching on constructing the Arrow RecordBatch; a batch will be |
| serialized to the Arrow RecordBatch when the total number of records exceeds the |
| configured threshold. |
| """ |
| # a set of variables for the state of current batch which will be converted to Arrow |
| # RecordBatch. |
| pdfs = [] |
| state_pdfs = [] |
| pdf_data_cnt = 0 |
| state_data_cnt = 0 |
| |
| return_schema = None |
| |
| for data in iterator: |
| # data represents the result of each call of user function |
| packaged_result = data[0] |
| |
| # There are two results from the call of user function: |
| # 1) iterator of pandas DataFrame (output) |
| # 2) updated state instance |
| pdf_iter = packaged_result[0][0] |
| state = packaged_result[0][1] |
| |
| # This is static and won't change across batches. |
| return_schema = packaged_result[1] |
| |
| for pdf in pdf_iter: |
| # We ignore empty pandas DataFrame. |
| if len(pdf) > 0: |
| pdf_data_cnt += len(pdf) |
| pdfs.append(pdf) |
| |
| # If the total number of records in current batch exceeds the configured |
| # threshold, time to construct the Arrow RecordBatch from the batch. |
| if pdf_data_cnt > self.arrow_max_records_per_batch: |
| batch = construct_record_batch( |
| pdfs, pdf_data_cnt, return_schema, state_pdfs, state_data_cnt |
| ) |
| |
| # Reset the variables to start with new batch for further data. |
| pdfs = [] |
| state_pdfs = [] |
| pdf_data_cnt = 0 |
| state_data_cnt = 0 |
| |
| yield batch |
| |
| # This has to be performed 'after' evaluating all elements in iterator, so that |
| # the user function has been completed and the state is guaranteed to be updated. |
| state_pdf = construct_state_pdf(state) |
| |
| state_pdfs.append(state_pdf) |
| state_data_cnt += 1 |
| |
| # processed all output, but current batch may not be flushed yet. |
| if pdf_data_cnt > 0 or state_data_cnt > 0: |
| batch = construct_record_batch( |
| pdfs, pdf_data_cnt, return_schema, state_pdfs, state_data_cnt |
| ) |
| |
| yield batch |
| |
| batches = self._write_stream_start(serialize_batches(), stream) |
| return ArrowStreamSerializer.dump_stream(self, batches, stream) |
| |
| |
| class TransformWithStateInPandasSerializer(ArrowStreamPandasUDFSerializer): |
| """ |
| Serializer used by Python worker to evaluate UDF for |
| :meth:`pyspark.sql.GroupedData.transformWithStateInPandasSerializer`. |
| |
| Parameters |
| ---------- |
| timezone : str |
| A timezone to respect when handling timestamp values |
| safecheck : bool |
| If True, conversion from Arrow to Pandas checks for overflow/truncation |
| assign_cols_by_name : bool |
| If True, then Pandas DataFrames will get columns by name |
| arrow_max_records_per_batch : int |
| Limit of the number of records that can be written to a single ArrowRecordBatch in memory. |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| prefer_int_ext_dtype, |
| arrow_max_records_per_batch, |
| arrow_max_bytes_per_batch, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| assign_cols_by_name=assign_cols_by_name, |
| df_for_struct=False, |
| struct_in_pandas="dict", |
| ndarray_as_list=False, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| arrow_cast=True, |
| input_type=None, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| ) |
| self.arrow_max_records_per_batch = ( |
| arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 |
| ) |
| self.arrow_max_bytes_per_batch = arrow_max_bytes_per_batch |
| self.key_offsets = None |
| self.average_arrow_row_size = 0 |
| self.total_bytes = 0 |
| self.total_rows = 0 |
| |
| def _update_batch_size_stats(self, batch): |
| """ |
| Update batch size statistics for adaptive batching. |
| """ |
| # Short circuit batch size calculation if the batch size is |
| # unlimited as computing batch size is computationally expensive. |
| if self.arrow_max_bytes_per_batch != 2**31 - 1 and batch.num_rows > 0: |
| batch_bytes = sum( |
| buf.size for col in batch.columns for buf in col.buffers() if buf is not None |
| ) |
| self.total_bytes += batch_bytes |
| self.total_rows += batch.num_rows |
| self.average_arrow_row_size = self.total_bytes / self.total_rows |
| |
| def load_stream(self, stream): |
| """ |
| Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunk, and |
| convert the data into Rows. |
| |
| Please refer the doc of inner function `generate_data_batches` for more details how |
| this function works in overall. |
| """ |
| import pandas as pd |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasFuncMode, |
| ) |
| |
| def generate_data_batches(batches): |
| """ |
| Deserialize ArrowRecordBatches and return a generator of Rows. |
| |
| The deserialization logic assumes that Arrow RecordBatches contain the data with the |
| ordering that data chunks for same grouping key will appear sequentially. |
| |
| This function must avoid materializing multiple Arrow RecordBatches into memory at the |
| same time. And data chunks from the same grouping key should appear sequentially. |
| """ |
| |
| def row_stream(): |
| for batch in batches: |
| self._update_batch_size_stats(batch) |
| data_pandas = ArrowBatchTransformer.to_pandas( |
| batch, |
| timezone=self._timezone, |
| schema=self._input_type, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ) |
| for row in pd.concat(data_pandas, axis=1).itertuples(index=False): |
| batch_key = tuple(row[s] for s in self.key_offsets) |
| yield (batch_key, row) |
| |
| for batch_key, group_rows in groupby(row_stream(), key=lambda x: x[0]): |
| rows = [] |
| for _, row in group_rows: |
| rows.append(row) |
| if ( |
| len(rows) >= self.arrow_max_records_per_batch |
| or len(rows) * self.average_arrow_row_size >= self.arrow_max_bytes_per_batch |
| ): |
| yield (batch_key, pd.DataFrame(rows)) |
| rows = [] |
| if rows: |
| yield (batch_key, pd.DataFrame(rows)) |
| |
| _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) |
| data_batches = generate_data_batches(_batches) |
| |
| for k, g in groupby(data_batches, key=lambda x: x[0]): |
| yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g) |
| |
| yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) |
| |
| yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Read through an iterator of (iterator of pandas DataFrame), serialize them to Arrow |
| RecordBatches, and write batches to stream. |
| """ |
| |
| def flatten_iterator(): |
| # iterator: iter[list[(iter[pandas.DataFrame], pdf_type)]] |
| for packed in iterator: |
| iter_pdf_with_type = packed[0] |
| iter_pdf = iter_pdf_with_type[0] |
| pdf_type = iter_pdf_with_type[1] |
| for pdf in iter_pdf: |
| yield [(pdf, pdf_type)] |
| |
| super().dump_stream(flatten_iterator(), stream) |
| |
| |
| class TransformWithStateInPandasInitStateSerializer(TransformWithStateInPandasSerializer): |
| """ |
| Serializer used by Python worker to evaluate UDF for |
| :meth:`pyspark.sql.GroupedData.transformWithStateInPandasInitStateSerializer`. |
| Parameters |
| ---------- |
| Same as input parameters in TransformWithStateInPandasSerializer. |
| """ |
| |
| def __init__( |
| self, |
| *, |
| timezone, |
| safecheck, |
| assign_cols_by_name, |
| prefer_int_ext_dtype, |
| arrow_max_records_per_batch, |
| arrow_max_bytes_per_batch, |
| int_to_decimal_coercion_enabled, |
| ): |
| super().__init__( |
| timezone=timezone, |
| safecheck=safecheck, |
| assign_cols_by_name=assign_cols_by_name, |
| prefer_int_ext_dtype=prefer_int_ext_dtype, |
| arrow_max_records_per_batch=arrow_max_records_per_batch, |
| arrow_max_bytes_per_batch=arrow_max_bytes_per_batch, |
| int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, |
| ) |
| self.init_key_offsets = None |
| |
| def load_stream(self, stream): |
| import pyarrow as pa |
| import pandas as pd |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasFuncMode, |
| ) |
| |
| def generate_data_batches(batches): |
| """ |
| Deserialize ArrowRecordBatches and return a generator of pandas.Series list. |
| |
| The deserialization logic assumes that Arrow RecordBatches contain the data with the |
| ordering that data chunks for same grouping key will appear sequentially. |
| See `TransformWithStateInPandasPythonInitialStateRunner` for arrow batch schema sent |
| from JVM. |
| This function flatten the columns of input rows and initial state rows and feed them |
| into the data generator. |
| """ |
| |
| def flatten_columns(cur_batch, col_name): |
| state_column = cur_batch.column(cur_batch.schema.get_field_index(col_name)) |
| |
| # Check if the entire column is null |
| if state_column.null_count == len(state_column): |
| # Return empty table with no columns |
| return pa.Table.from_arrays([], names=[]) |
| |
| state_field_names = [ |
| state_column.type[i].name for i in range(state_column.type.num_fields) |
| ] |
| state_field_arrays = [ |
| state_column.field(i) for i in range(state_column.type.num_fields) |
| ] |
| table_from_fields = pa.Table.from_arrays( |
| state_field_arrays, names=state_field_names |
| ) |
| return table_from_fields |
| |
| """ |
| The arrow batch is written in the schema: |
| schema: StructType = new StructType() |
| .add("inputData", dataSchema) |
| .add("initState", initStateSchema) |
| We'll parse batch into Tuples of (key, inputData, initState) and pass into the Python |
| data generator. Rows in the same batch may have different grouping keys, |
| but each batch will have either init_data or input_data, not mix. |
| """ |
| |
| def to_pandas(table): |
| return ArrowBatchTransformer.to_pandas( |
| table, |
| timezone=self._timezone, |
| schema=self._input_type, |
| struct_in_pandas=self._struct_in_pandas, |
| ndarray_as_list=self._ndarray_as_list, |
| prefer_int_ext_dtype=self._prefer_int_ext_dtype, |
| df_for_struct=self._df_for_struct, |
| ) |
| |
| def row_stream(): |
| for batch in batches: |
| self._update_batch_size_stats(batch) |
| |
| data_table = flatten_columns(batch, "inputData") |
| init_table = flatten_columns(batch, "initState") |
| |
| # Check column count - empty table has no columns |
| has_data = data_table.num_columns > 0 |
| has_init = init_table.num_columns > 0 |
| |
| assert not (has_data and has_init) |
| |
| if has_data: |
| for row in pd.concat(to_pandas(data_table), axis=1).itertuples(index=False): |
| batch_key = tuple(row[s] for s in self.key_offsets) |
| yield (batch_key, row, None) |
| elif has_init: |
| for row in pd.concat(to_pandas(init_table), axis=1).itertuples(index=False): |
| batch_key = tuple(row[s] for s in self.init_key_offsets) |
| yield (batch_key, None, row) |
| |
| EMPTY_DATAFRAME = pd.DataFrame() |
| for batch_key, group_rows in groupby(row_stream(), key=lambda x: x[0]): |
| rows = [] |
| init_state_rows = [] |
| for _, row, init_state_row in group_rows: |
| if row is not None: |
| rows.append(row) |
| if init_state_row is not None: |
| init_state_rows.append(init_state_row) |
| |
| total_len = len(rows) + len(init_state_rows) |
| if ( |
| total_len >= self.arrow_max_records_per_batch |
| or total_len * self.average_arrow_row_size >= self.arrow_max_bytes_per_batch |
| ): |
| yield ( |
| batch_key, |
| pd.DataFrame(rows) if len(rows) > 0 else EMPTY_DATAFRAME.copy(), |
| ( |
| pd.DataFrame(init_state_rows) |
| if len(init_state_rows) > 0 |
| else EMPTY_DATAFRAME.copy() |
| ), |
| ) |
| rows = [] |
| init_state_rows = [] |
| if rows or init_state_rows: |
| yield ( |
| batch_key, |
| pd.DataFrame(rows) if len(rows) > 0 else EMPTY_DATAFRAME.copy(), |
| ( |
| pd.DataFrame(init_state_rows) |
| if len(init_state_rows) > 0 |
| else EMPTY_DATAFRAME.copy() |
| ), |
| ) |
| |
| _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) |
| data_batches = generate_data_batches(_batches) |
| |
| for k, g in groupby(data_batches, key=lambda x: x[0]): |
| yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g) |
| |
| yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) |
| |
| yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) |
| |
| |
| class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer): |
| """ |
| Serializer used by Python worker to evaluate UDF for |
| :meth:`pyspark.sql.GroupedData.transformWithState`. |
| |
| Parameters |
| ---------- |
| arrow_max_records_per_batch : int |
| Limit of the number of records that can be written to a single ArrowRecordBatch in memory. |
| """ |
| |
| def __init__(self, *, arrow_max_records_per_batch): |
| super().__init__() |
| self.arrow_max_records_per_batch = ( |
| arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 |
| ) |
| self.key_offsets = None |
| |
| def load_stream(self, stream): |
| """ |
| Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunks, |
| and convert the data into a list of pandas.Series. |
| |
| Please refer the doc of inner function `generate_data_batches` for more details how |
| this function works in overall. |
| """ |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasFuncMode, |
| ) |
| import itertools |
| |
| def generate_data_batches(batches): |
| """ |
| Deserialize ArrowRecordBatches and return a generator of Row. |
| |
| The deserialization logic assumes that Arrow RecordBatches contain the data with the |
| ordering that data chunks for same grouping key will appear sequentially. |
| |
| This function must avoid materializing multiple Arrow RecordBatches into memory at the |
| same time. And data chunks from the same grouping key should appear sequentially. |
| """ |
| for batch in batches: |
| DataRow = Row(*batch.schema.names) |
| |
| # Iterate row by row without converting the whole batch |
| num_cols = batch.num_columns |
| for row_idx in range(batch.num_rows): |
| # build the key for this row |
| row_key = tuple(batch[o][row_idx].as_py() for o in self.key_offsets) |
| row = DataRow(*(batch.column(i)[row_idx].as_py() for i in range(num_cols))) |
| yield row_key, row |
| |
| _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) |
| data_batches = generate_data_batches(_batches) |
| |
| for k, g in groupby(data_batches, key=lambda x: x[0]): |
| chained = itertools.chain(g) |
| chained_values = map(lambda x: x[1], chained) |
| yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, chained_values) |
| |
| yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) |
| |
| yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) |
| |
| def dump_stream(self, iterator, stream): |
| """ |
| Read through an iterator of (iterator of Row), serialize them to Arrow |
| RecordBatches, and write batches to stream. |
| """ |
| import pyarrow as pa |
| |
| from pyspark.sql.pandas.types import to_arrow_type |
| |
| def flatten_iterator(): |
| # iterator: iter[list[(iter[Row], spark_type)]] |
| for packed in iterator: |
| iter_row_with_type = packed[0] |
| iter_row = iter_row_with_type[0] |
| spark_type = iter_row_with_type[1] |
| |
| # Convert spark type to arrow type |
| # TODO: WE need to make this configurable, currently using default values. |
| arrow_type = to_arrow_type( |
| spark_type, |
| timezone="UTC", |
| prefers_large_types=False, |
| ) |
| |
| rows_as_dict = [] |
| for row in iter_row: |
| row_as_dict = row.asDict(True) |
| rows_as_dict.append(row_as_dict) |
| |
| pdf_schema = pa.schema(list(arrow_type)) |
| record_batch = pa.RecordBatch.from_pylist(rows_as_dict, schema=pdf_schema) |
| |
| yield (record_batch, arrow_type) |
| |
| return ArrowStreamUDFSerializer.dump_stream(self, flatten_iterator(), stream) |
| |
| |
| class TransformWithStateInPySparkRowInitStateSerializer(TransformWithStateInPySparkRowSerializer): |
| """ |
| Serializer used by Python worker to evaluate UDF for |
| :meth:`pyspark.sql.GroupedData.transformWithStateInPySparkRowInitStateSerializer`. |
| Parameters |
| ---------- |
| Same as input parameters in TransformWithStateInPySparkRowSerializer. |
| """ |
| |
| def __init__(self, *, arrow_max_records_per_batch): |
| super().__init__(arrow_max_records_per_batch=arrow_max_records_per_batch) |
| self.init_key_offsets = None |
| |
| def load_stream(self, stream): |
| import pyarrow as pa |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasFuncMode, |
| ) |
| |
| def generate_data_batches(batches) -> Iterator[Tuple[Any, Optional[Any], Optional[Any]]]: |
| """ |
| Deserialize ArrowRecordBatches and return a generator of Row. |
| The deserialization logic assumes that Arrow RecordBatches contain the data with the |
| ordering that data chunks for same grouping key will appear sequentially. |
| See `TransformWithStateInPySparkPythonInitialStateRunner` for arrow batch schema sent |
| from JVM. |
| This function flattens the columns of input rows and initial state rows and feed them |
| into the data generator. |
| """ |
| |
| def extract_rows( |
| cur_batch, col_name, key_offsets |
| ) -> Optional[Iterator[Tuple[Any, Any]]]: |
| data_column = cur_batch.column(cur_batch.schema.get_field_index(col_name)) |
| |
| # Check if the entire column is null |
| if data_column.null_count == len(data_column): |
| return None |
| |
| data_field_names = [ |
| data_column.type[i].name for i in range(data_column.type.num_fields) |
| ] |
| data_field_arrays = [ |
| data_column.field(i) for i in range(data_column.type.num_fields) |
| ] |
| |
| DataRow = Row(*data_field_names) |
| |
| table = pa.Table.from_arrays(data_field_arrays, names=data_field_names) |
| |
| if table.num_rows == 0: |
| return None |
| |
| def row_iterator(): |
| for row_idx in range(table.num_rows): |
| key = tuple(table.column(o)[row_idx].as_py() for o in key_offsets) |
| row = DataRow( |
| *(table.column(i)[row_idx].as_py() for i in range(table.num_columns)) |
| ) |
| yield (key, row) |
| |
| return row_iterator() |
| |
| """ |
| The arrow batch is written in the schema: |
| schema: StructType = new StructType() |
| .add("inputData", dataSchema) |
| .add("initState", initStateSchema) |
| We'll parse batch into Tuples of (key, inputData, initState) and pass into the Python |
| data generator. Each batch will have either init_data or input_data, not mix. |
| """ |
| for batch in batches: |
| # Detect which column has data - each batch contains only one type |
| input_result = extract_rows(batch, "inputData", self.key_offsets) |
| init_result = extract_rows(batch, "initState", self.init_key_offsets) |
| |
| assert not (input_result is not None and init_result is not None) |
| |
| if input_result is not None: |
| for key, input_data_row in input_result: |
| yield (key, input_data_row, None) |
| elif init_result is not None: |
| for key, init_state_row in init_result: |
| yield (key, None, init_state_row) |
| |
| _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) |
| data_batches = generate_data_batches(_batches) |
| |
| for k, g in groupby(data_batches, key=lambda x: x[0]): |
| input_rows = [] |
| init_rows = [] |
| |
| for batch_key, input_row, init_row in g: |
| if input_row is not None: |
| input_rows.append(input_row) |
| if init_row is not None: |
| init_rows.append(init_row) |
| |
| total_len = len(input_rows) + len(init_rows) |
| if total_len >= self.arrow_max_records_per_batch: |
| ret_tuple = (iter(input_rows), iter(init_rows)) |
| yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple) |
| input_rows = [] |
| init_rows = [] |
| |
| if input_rows or init_rows: |
| ret_tuple = (iter(input_rows), iter(init_rows)) |
| yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple) |
| |
| yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) |
| |
| yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) |
