| # |
| # 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. |
| # |
| from __future__ import annotations |
| |
| import inspect |
| import uuid |
| from typing import Any, Callable, Iterator, List, Mapping, TYPE_CHECKING, Tuple, Union, Optional |
| |
| import numpy as np |
| |
| try: |
| import pandas as pd |
| except ImportError: |
| pass # Let it throw a better error message later when the API is invoked. |
| |
| from pyspark.sql.functions import pandas_udf |
| from pyspark.sql.column import Column |
| from pyspark.sql.types import ( |
| ArrayType, |
| ByteType, |
| DataType, |
| DoubleType, |
| FloatType, |
| IntegerType, |
| LongType, |
| ShortType, |
| StringType, |
| StructType, |
| ) |
| from pyspark.ml.util import try_remote_functions |
| |
| if TYPE_CHECKING: |
| from pyspark.sql._typing import UserDefinedFunctionLike |
| |
| supported_scalar_types = ( |
| ByteType, |
| ShortType, |
| IntegerType, |
| LongType, |
| FloatType, |
| DoubleType, |
| StringType, |
| ) |
| |
| # Callable type for end user predict functions that take a variable number of ndarrays as |
| # input and returns one of the following as output: |
| # - single ndarray (single output) |
| # - dictionary of named ndarrays (multiple outputs represented in columnar form) |
| # - list of dictionaries of named ndarrays (multiple outputs represented in row form) |
| PredictBatchFunction = Callable[ |
| [np.ndarray], Union[np.ndarray, Mapping[str, np.ndarray], List[Mapping[str, np.dtype]]] |
| ] |
| |
| |
| @try_remote_functions |
| def vector_to_array(col: Column, dtype: str = "float64") -> Column: |
| """ |
| Converts a column of MLlib sparse/dense vectors into a column of dense arrays. |
| |
| .. versionadded:: 3.0.0 |
| |
| .. versionchanged:: 3.5.0 |
| Supports Spark Connect. |
| |
| Parameters |
| ---------- |
| col : :py:class:`pyspark.sql.Column` or str |
| Input column |
| dtype : str, optional |
| The data type of the output array. Valid values: "float64" or "float32". |
| |
| Returns |
| ------- |
| :py:class:`pyspark.sql.Column` |
| The converted column of dense arrays. |
| |
| Examples |
| -------- |
| >>> from pyspark.ml.linalg import Vectors |
| >>> from pyspark.ml.functions import vector_to_array |
| >>> from pyspark.mllib.linalg import Vectors as OldVectors |
| >>> df = spark.createDataFrame([ |
| ... (Vectors.dense(1.0, 2.0, 3.0), OldVectors.dense(10.0, 20.0, 30.0)), |
| ... (Vectors.sparse(3, [(0, 2.0), (2, 3.0)]), |
| ... OldVectors.sparse(3, [(0, 20.0), (2, 30.0)]))], |
| ... ["vec", "oldVec"]) |
| >>> df1 = df.select(vector_to_array("vec").alias("vec"), |
| ... vector_to_array("oldVec").alias("oldVec")) |
| >>> df1.collect() |
| [Row(vec=[1.0, 2.0, 3.0], oldVec=[10.0, 20.0, 30.0]), |
| Row(vec=[2.0, 0.0, 3.0], oldVec=[20.0, 0.0, 30.0])] |
| >>> df2 = df.select(vector_to_array("vec", "float32").alias("vec"), |
| ... vector_to_array("oldVec", "float32").alias("oldVec")) |
| >>> df2.collect() |
| [Row(vec=[1.0, 2.0, 3.0], oldVec=[10.0, 20.0, 30.0]), |
| Row(vec=[2.0, 0.0, 3.0], oldVec=[20.0, 0.0, 30.0])] |
| >>> df1.schema.fields |
| [StructField('vec', ArrayType(DoubleType(), False), False), |
| StructField('oldVec', ArrayType(DoubleType(), False), False)] |
| >>> df2.schema.fields |
| [StructField('vec', ArrayType(FloatType(), False), False), |
| StructField('oldVec', ArrayType(FloatType(), False), False)] |
| """ |
| from pyspark.core.context import SparkContext |
| from pyspark.sql.classic.column import Column, _to_java_column |
| |
| sc = SparkContext._active_spark_context |
| assert sc is not None and sc._jvm is not None |
| return Column( |
| getattr(sc._jvm, "org.apache.spark.ml.functions").vector_to_array( |
| _to_java_column(col), dtype |
| ) |
| ) |
| |
| |
| @try_remote_functions |
| def array_to_vector(col: Column) -> Column: |
| """ |
| Converts a column of array of numeric type into a column of pyspark.ml.linalg.DenseVector |
| instances |
| |
| .. versionadded:: 3.1.0 |
| |
| .. versionchanged:: 3.5.0 |
| Supports Spark Connect. |
| |
| Parameters |
| ---------- |
| col : :py:class:`pyspark.sql.Column` or str |
| Input column |
| |
| Returns |
| ------- |
| :py:class:`pyspark.sql.Column` |
| The converted column of dense vectors. |
| |
| Examples |
| -------- |
| >>> from pyspark.ml.functions import array_to_vector |
| >>> df1 = spark.createDataFrame([([1.5, 2.5],),], schema='v1 array<double>') |
| >>> df1.select(array_to_vector('v1').alias('vec1')).collect() |
| [Row(vec1=DenseVector([1.5, 2.5]))] |
| >>> df2 = spark.createDataFrame([([1.5, 3.5],),], schema='v1 array<float>') |
| >>> df2.select(array_to_vector('v1').alias('vec1')).collect() |
| [Row(vec1=DenseVector([1.5, 3.5]))] |
| >>> df3 = spark.createDataFrame([([1, 3],),], schema='v1 array<int>') |
| >>> df3.select(array_to_vector('v1').alias('vec1')).collect() |
| [Row(vec1=DenseVector([1.0, 3.0]))] |
| """ |
| from pyspark.core.context import SparkContext |
| from pyspark.sql.classic.column import Column, _to_java_column |
| |
| sc = SparkContext._active_spark_context |
| assert sc is not None and sc._jvm is not None |
| return Column( |
| getattr(sc._jvm, "org.apache.spark.ml.functions").array_to_vector(_to_java_column(col)) |
| ) |
| |
| |
| def _batched( |
| data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]], batch_size: int |
| ) -> Iterator[pd.DataFrame]: |
| """Generator that splits a pandas dataframe/series into batches.""" |
| if isinstance(data, pd.DataFrame): |
| df = data |
| elif isinstance(data, pd.Series): |
| df = pd.concat((data,), axis=1) |
| else: # isinstance(data, Tuple[pd.Series]): |
| df = pd.concat(data, axis=1) |
| |
| index = 0 |
| data_size = len(df) |
| while index < data_size: |
| yield df.iloc[index : index + batch_size] |
| index += batch_size |
| |
| |
| def _is_tensor_col(data: Union[pd.Series, pd.DataFrame]) -> bool: |
| if isinstance(data, pd.Series): |
| return data.dtype == np.object_ and isinstance(data.iloc[0], (np.ndarray, list)) |
| elif isinstance(data, pd.DataFrame): |
| return any(data.dtypes == np.object_) and any( |
| [isinstance(d, (np.ndarray, list)) for d in data.iloc[0]] |
| ) |
| else: |
| raise ValueError( |
| "Unexpected data type: {}, expected pd.Series or pd.DataFrame.".format(type(data)) |
| ) |
| |
| |
| def _has_tensor_cols(data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]]) -> bool: |
| """Check if input Series/DataFrame/Tuple contains any tensor-valued columns.""" |
| if isinstance(data, (pd.Series, pd.DataFrame)): |
| return _is_tensor_col(data) |
| else: # isinstance(data, Tuple): |
| return any(_is_tensor_col(elem) for elem in data) |
| |
| |
| def _validate_and_transform_multiple_inputs( |
| batch: pd.DataFrame, input_shapes: List[Optional[List[int]]], num_input_cols: int |
| ) -> List[np.ndarray]: |
| multi_inputs = [batch[col].to_numpy() for col in batch.columns] |
| if input_shapes: |
| if len(input_shapes) == num_input_cols: |
| multi_inputs = [ |
| np.vstack(v).reshape([-1] + input_shapes[i]) # type: ignore |
| if input_shapes[i] |
| else v |
| for i, v in enumerate(multi_inputs) |
| ] |
| if not all([len(x) == len(batch) for x in multi_inputs]): |
| raise ValueError("Input data does not match expected shape.") |
| else: |
| raise ValueError("input_tensor_shapes must match columns") |
| |
| return multi_inputs |
| |
| |
| def _validate_and_transform_single_input( |
| batch: pd.DataFrame, |
| input_shapes: List[List[int] | None], |
| has_tensors: bool, |
| has_tuple: bool, |
| ) -> np.ndarray: |
| # multiple input columns for single expected input |
| if has_tensors: |
| # tensor columns |
| if len(batch.columns) == 1: |
| # one tensor column and one expected input, vstack rows |
| single_input = np.vstack(batch.iloc[:, 0]) |
| else: |
| raise ValueError( |
| "Multiple input columns found, but model expected a single " |
| "input, use `array` to combine columns into tensors." |
| ) |
| else: |
| # scalar columns |
| if len(batch.columns) == 1: |
| # single scalar column, remove extra dim |
| np_batch = batch.to_numpy() |
| single_input = np.squeeze(np_batch, -1) if len(np_batch.shape) > 1 else np_batch |
| if input_shapes and input_shapes[0] not in [None, [], [1]]: |
| raise ValueError("Invalid input_tensor_shape for scalar column.") |
| elif not has_tuple: |
| # columns grouped via `array`, convert to single tensor |
| single_input = batch.to_numpy() |
| if input_shapes and input_shapes[0] != [len(batch.columns)]: |
| raise ValueError("Input data does not match expected shape.") |
| else: |
| raise ValueError( |
| "Multiple input columns found, but model expected a single " |
| "input, use `array` to combine columns into tensors." |
| ) |
| |
| # if input_tensor_shapes provided, try to reshape input |
| if input_shapes: |
| if len(input_shapes) == 1: |
| single_input = single_input.reshape([-1] + input_shapes[0]) # type: ignore |
| if len(single_input) != len(batch): |
| raise ValueError("Input data does not match expected shape.") |
| else: |
| raise ValueError("Multiple input_tensor_shapes found, but model expected one input") |
| |
| return single_input |
| |
| |
| def _validate_and_transform_prediction_result( |
| preds: np.ndarray | Mapping[str, np.ndarray] | List[Mapping[str, Any]], |
| num_input_rows: int, |
| return_type: DataType, |
| ) -> pd.DataFrame | pd.Series: |
| """Validate numpy-based model predictions against the expected pandas_udf return_type and |
| transforms the predictions into an equivalent pandas DataFrame or Series.""" |
| if isinstance(return_type, StructType): |
| struct_rtype: StructType = return_type |
| fieldNames = struct_rtype.names |
| if isinstance(preds, dict): |
| # dictionary of columns |
| predNames = list(preds.keys()) |
| for field in struct_rtype.fields: |
| if isinstance(field.dataType, ArrayType): |
| if len(preds[field.name].shape) == 2: |
| preds[field.name] = list(preds[field.name]) |
| else: |
| raise ValueError( |
| "Prediction results for ArrayType must be two-dimensional." |
| ) |
| elif isinstance(field.dataType, supported_scalar_types): |
| if len(preds[field.name].shape) != 1: |
| raise ValueError( |
| "Prediction results for scalar types must be one-dimensional." |
| ) |
| else: |
| raise ValueError("Unsupported field type in return struct type.") |
| |
| if len(preds[field.name]) != num_input_rows: |
| raise ValueError("Prediction results must have same length as input data") |
| |
| elif isinstance(preds, list) and isinstance(preds[0], dict): |
| # rows of dictionaries |
| predNames = list(preds[0].keys()) |
| if len(preds) != num_input_rows: |
| raise ValueError("Prediction results must have same length as input data.") |
| for field in struct_rtype.fields: |
| if isinstance(field.dataType, ArrayType): |
| if len(preds[0][field.name].shape) != 1: |
| raise ValueError( |
| "Prediction results for ArrayType must be one-dimensional." |
| ) |
| elif isinstance(field.dataType, supported_scalar_types): |
| if not np.isscalar(preds[0][field.name]): |
| raise ValueError("Invalid scalar prediction result.") |
| else: |
| raise ValueError("Unsupported field type in return struct type.") |
| else: |
| raise ValueError( |
| "Prediction results for StructType must be a dictionary or " |
| "a list of dictionary, got: {}".format(type(preds)) |
| ) |
| |
| # check column names |
| if set(predNames) != set(fieldNames): |
| raise ValueError( |
| "Prediction result columns did not match expected return_type " |
| "columns: expected {}, got: {}".format(fieldNames, predNames) |
| ) |
| |
| return pd.DataFrame(preds) |
| elif isinstance(return_type, ArrayType): |
| if isinstance(preds, np.ndarray): |
| if len(preds) != num_input_rows: |
| raise ValueError("Prediction results must have same length as input data.") |
| if len(preds.shape) != 2: |
| raise ValueError("Prediction results for ArrayType must be two-dimensional.") |
| else: |
| raise ValueError("Prediction results for ArrayType must be an ndarray.") |
| |
| return pd.Series(list(preds)) |
| elif isinstance(return_type, supported_scalar_types): |
| preds_array: np.ndarray = preds # type: ignore |
| if len(preds_array) != num_input_rows: |
| raise ValueError("Prediction results must have same length as input data.") |
| if not ( |
| (len(preds_array.shape) == 2 and preds_array.shape[1] == 1) |
| or len(preds_array.shape) == 1 |
| ): |
| raise ValueError("Invalid shape for scalar prediction result.") |
| |
| output = np.squeeze(preds_array, -1) if len(preds_array.shape) > 1 else preds_array |
| return pd.Series(output).astype(output.dtype) |
| else: |
| raise ValueError("Unsupported return type") |
| |
| |
| def predict_batch_udf( |
| make_predict_fn: Callable[ |
| [], |
| PredictBatchFunction, |
| ], |
| *, |
| return_type: DataType, |
| batch_size: int, |
| input_tensor_shapes: Optional[Union[List[Optional[List[int]]], Mapping[int, List[int]]]] = None, |
| ) -> UserDefinedFunctionLike: |
| """Given a function which loads a model and returns a `predict` function for inference over a |
| batch of numpy inputs, returns a Pandas UDF wrapper for inference over a Spark DataFrame. |
| |
| The returned Pandas UDF does the following on each DataFrame partition: |
| |
| * calls the `make_predict_fn` to load the model and cache its `predict` function. |
| * batches the input records as numpy arrays and invokes `predict` on each batch. |
| |
| Note: this assumes that the `make_predict_fn` encapsulates all of the necessary dependencies for |
| running the model, or the Spark executor environment already satisfies all runtime requirements. |
| |
| For the conversion of the Spark DataFrame to numpy arrays, there is a one-to-one mapping between |
| the input arguments of the `predict` function (returned by the `make_predict_fn`) and the input |
| columns sent to the Pandas UDF (returned by the `predict_batch_udf`) at runtime. Each input |
| column will be converted as follows: |
| |
| * scalar column -> 1-dim np.ndarray |
| * tensor column + tensor shape -> N-dim np.ndarray |
| |
| Note that any tensor columns in the Spark DataFrame must be represented as a flattened |
| one-dimensional array, and multiple scalar columns can be combined into a single tensor column |
| using the standard :py:func:`pyspark.sql.functions.array()` function. |
| |
| .. versionadded:: 3.4.0 |
| |
| Parameters |
| ---------- |
| make_predict_fn : callable |
| Function which is responsible for loading a model and returning a |
| :py:class:`PredictBatchFunction` which takes one or more numpy arrays as input and returns |
| one of the following: |
| |
| * a numpy array (for a single output) |
| * a dictionary of named numpy arrays (for multiple outputs) |
| * a row-oriented list of dictionaries (for multiple outputs). |
| |
| For a dictionary of named numpy arrays, the arrays can only be one or two dimensional, since |
| higher dimensional arrays are not supported. For a row-oriented list of dictionaries, each |
| element in the dictionary must be either a scalar or one-dimensional array. |
| return_type : :py:class:`pyspark.sql.types.DataType` or str. |
| Spark SQL datatype for the expected output: |
| |
| * Scalar (e.g. IntegerType, FloatType) --> 1-dim numpy array. |
| * ArrayType --> 2-dim numpy array. |
| * StructType --> dict with keys matching struct fields. |
| * StructType --> list of dict with keys matching struct fields, for models like the |
| `Huggingface pipeline for sentiment analysis |
| <https://huggingface.co/docs/transformers/quicktour#pipeline-usage>`_. |
| |
| batch_size : int |
| Batch size to use for inference. This is typically a limitation of the model |
| and/or available hardware resources and is usually smaller than the Spark partition size. |
| input_tensor_shapes : list, dict, optional. |
| A list of ints or a dictionary of ints (key) and list of ints (value). |
| Input tensor shapes for models with tensor inputs. This can be a list of shapes, |
| where each shape is a list of integers or None (for scalar inputs). Alternatively, this |
| can be represented by a "sparse" dictionary, where the keys are the integer indices of the |
| inputs, and the values are the shapes. Each tensor input value in the Spark DataFrame must |
| be represented as a single column containing a flattened 1-D array. The provided |
| `input_tensor_shapes` will be used to reshape the flattened array into the expected tensor |
| shape. For the list form, the order of the tensor shapes must match the order of the |
| selected DataFrame columns. The batch dimension (typically -1 or None in the first |
| dimension) should not be included, since it will be determined by the batch_size argument. |
| Tabular datasets with scalar-valued columns should not provide this argument. |
| |
| Returns |
| ------- |
| :py:class:`UserDefinedFunctionLike` |
| A Pandas UDF for model inference on a Spark DataFrame. |
| |
| Examples |
| -------- |
| For a pre-trained TensorFlow MNIST model with two-dimensional input images represented as a |
| flattened tensor value stored in a single Spark DataFrame column of type `array<float>`. |
| |
| .. code-block:: python |
| |
| from pyspark.ml.functions import predict_batch_udf |
| |
| def make_mnist_fn(): |
| # load/init happens once per python worker |
| import tensorflow as tf |
| model = tf.keras.models.load_model('/path/to/mnist_model') |
| |
| # predict on batches of tasks/partitions, using cached model |
| def predict(inputs: np.ndarray) -> np.ndarray: |
| # inputs.shape = [batch_size, 784], see input_tensor_shapes |
| # outputs.shape = [batch_size, 10], see return_type |
| return model.predict(inputs) |
| |
| return predict |
| |
| mnist_udf = predict_batch_udf(make_mnist_fn, |
| return_type=ArrayType(FloatType()), |
| batch_size=100, |
| input_tensor_shapes=[[784]]) |
| |
| df = spark.read.parquet("/path/to/mnist_data") |
| df.show(5) |
| # +--------------------+ |
| # | data| |
| # +--------------------+ |
| # |[0.0, 0.0, 0.0, 0...| |
| # |[0.0, 0.0, 0.0, 0...| |
| # |[0.0, 0.0, 0.0, 0...| |
| # |[0.0, 0.0, 0.0, 0...| |
| # |[0.0, 0.0, 0.0, 0...| |
| # +--------------------+ |
| |
| df.withColumn("preds", mnist_udf("data")).show(5) |
| # +--------------------+--------------------+ |
| # | data| preds| |
| # +--------------------+--------------------+ |
| # |[0.0, 0.0, 0.0, 0...|[-13.511008, 8.84...| |
| # |[0.0, 0.0, 0.0, 0...|[-5.3957458, -2.2...| |
| # |[0.0, 0.0, 0.0, 0...|[-7.2014456, -8.8...| |
| # |[0.0, 0.0, 0.0, 0...|[-19.466187, -13....| |
| # |[0.0, 0.0, 0.0, 0...|[-5.7757926, -7.8...| |
| # +--------------------+--------------------+ |
| |
| To demonstrate usage with different combinations of input and output types, the following |
| examples just use simple mathematical transforms as the models. |
| |
| * Single scalar column |
| Input DataFrame has a single scalar column, which will be passed to the `predict` |
| function as a 1-D numpy array. |
| |
| >>> import numpy as np |
| >>> import pandas as pd |
| >>> from pyspark.ml.functions import predict_batch_udf |
| >>> from pyspark.sql.types import FloatType |
| >>> |
| >>> df = spark.createDataFrame(pd.DataFrame(np.arange(100))) |
| >>> df.show(5) |
| +---+ |
| | 0| |
| +---+ |
| | 0| |
| | 1| |
| | 2| |
| | 3| |
| | 4| |
| +---+ |
| only showing top 5 rows |
| |
| >>> def make_times_two_fn(): |
| ... def predict(inputs: np.ndarray) -> np.ndarray: |
| ... # inputs.shape = [batch_size] |
| ... # outputs.shape = [batch_size] |
| ... return inputs * 2 |
| ... return predict |
| ... |
| >>> times_two_udf = predict_batch_udf(make_times_two_fn, |
| ... return_type=FloatType(), |
| ... batch_size=10) |
| >>> df = spark.createDataFrame(pd.DataFrame(np.arange(100))) |
| >>> df.withColumn("x2", times_two_udf("0")).show(5) |
| +---+---+ |
| | 0| x2| |
| +---+---+ |
| | 0|0.0| |
| | 1|2.0| |
| | 2|4.0| |
| | 3|6.0| |
| | 4|8.0| |
| +---+---+ |
| only showing top 5 rows |
| |
| * Multiple scalar columns |
| Input DataFrame has multiple columns of scalar values. If the user-provided `predict` |
| function expects a single input, then the user must combine the multiple columns into a |
| single tensor using `pyspark.sql.functions.array`. |
| |
| >>> import numpy as np |
| >>> import pandas as pd |
| >>> from pyspark.ml.functions import predict_batch_udf |
| >>> from pyspark.sql.functions import array |
| >>> |
| >>> data = np.arange(0, 1000, dtype=np.float64).reshape(-1, 4) |
| >>> pdf = pd.DataFrame(data, columns=['a','b','c','d']) |
| >>> df = spark.createDataFrame(pdf) |
| >>> df.show(5) |
| +----+----+----+----+ |
| | a| b| c| d| |
| +----+----+----+----+ |
| | 0.0| 1.0| 2.0| 3.0| |
| | 4.0| 5.0| 6.0| 7.0| |
| | 8.0| 9.0|10.0|11.0| |
| |12.0|13.0|14.0|15.0| |
| |16.0|17.0|18.0|19.0| |
| +----+----+----+----+ |
| only showing top 5 rows |
| |
| >>> def make_sum_fn(): |
| ... def predict(inputs: np.ndarray) -> np.ndarray: |
| ... # inputs.shape = [batch_size, 4] |
| ... # outputs.shape = [batch_size] |
| ... return np.sum(inputs, axis=1) |
| ... return predict |
| ... |
| >>> sum_udf = predict_batch_udf(make_sum_fn, |
| ... return_type=FloatType(), |
| ... batch_size=10, |
| ... input_tensor_shapes=[[4]]) |
| >>> df.withColumn("sum", sum_udf(array("a", "b", "c", "d"))).show(5) |
| +----+----+----+----+----+ |
| | a| b| c| d| sum| |
| +----+----+----+----+----+ |
| | 0.0| 1.0| 2.0| 3.0| 6.0| |
| | 4.0| 5.0| 6.0| 7.0|22.0| |
| | 8.0| 9.0|10.0|11.0|38.0| |
| |12.0|13.0|14.0|15.0|54.0| |
| |16.0|17.0|18.0|19.0|70.0| |
| +----+----+----+----+----+ |
| only showing top 5 rows |
| |
| If the `predict` function expects multiple inputs, then the number of selected input columns |
| must match the number of expected inputs. |
| |
| >>> def make_sum_fn(): |
| ... def predict(x1: np.ndarray, |
| ... x2: np.ndarray, |
| ... x3: np.ndarray, |
| ... x4: np.ndarray) -> np.ndarray: |
| ... # xN.shape = [batch_size] |
| ... # outputs.shape = [batch_size] |
| ... return x1 + x2 + x3 + x4 |
| ... return predict |
| ... |
| >>> sum_udf = predict_batch_udf(make_sum_fn, |
| ... return_type=FloatType(), |
| ... batch_size=10) |
| >>> df.withColumn("sum", sum_udf("a", "b", "c", "d")).show(5) |
| +----+----+----+----+----+ |
| | a| b| c| d| sum| |
| +----+----+----+----+----+ |
| | 0.0| 1.0| 2.0| 3.0| 6.0| |
| | 4.0| 5.0| 6.0| 7.0|22.0| |
| | 8.0| 9.0|10.0|11.0|38.0| |
| |12.0|13.0|14.0|15.0|54.0| |
| |16.0|17.0|18.0|19.0|70.0| |
| +----+----+----+----+----+ |
| only showing top 5 rows |
| |
| * Multiple tensor columns |
| Input DataFrame has multiple columns, where each column is a tensor. The number of columns |
| should match the number of expected inputs for the user-provided `predict` function. |
| |
| >>> import numpy as np |
| >>> import pandas as pd |
| >>> from pyspark.ml.functions import predict_batch_udf |
| >>> from pyspark.sql.types import ArrayType, FloatType, StructType, StructField |
| >>> from typing import Mapping |
| >>> |
| >>> data = np.arange(0, 1000, dtype=np.float64).reshape(-1, 4) |
| >>> pdf = pd.DataFrame(data, columns=['a','b','c','d']) |
| >>> pdf_tensor = pd.DataFrame() |
| >>> pdf_tensor['t1'] = pdf.values.tolist() |
| >>> pdf_tensor['t2'] = pdf.drop(columns='d').values.tolist() |
| >>> df = spark.createDataFrame(pdf_tensor) |
| >>> df.show(5) |
| +--------------------+------------------+ |
| | t1| t2| |
| +--------------------+------------------+ |
| |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| |
| |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]| |
| |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]| |
| |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]| |
| |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]| |
| +--------------------+------------------+ |
| only showing top 5 rows |
| |
| >>> def make_multi_sum_fn(): |
| ... def predict(x1: np.ndarray, x2: np.ndarray) -> np.ndarray: |
| ... # x1.shape = [batch_size, 4] |
| ... # x2.shape = [batch_size, 3] |
| ... # outputs.shape = [batch_size] |
| ... return np.sum(x1, axis=1) + np.sum(x2, axis=1) |
| ... return predict |
| ... |
| >>> multi_sum_udf = predict_batch_udf( |
| ... make_multi_sum_fn, |
| ... return_type=FloatType(), |
| ... batch_size=5, |
| ... input_tensor_shapes=[[4], [3]], |
| ... ) |
| >>> df.withColumn("sum", multi_sum_udf("t1", "t2")).show(5) |
| +--------------------+------------------+-----+ |
| | t1| t2| sum| |
| +--------------------+------------------+-----+ |
| |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 9.0| |
| |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]| 37.0| |
| |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]| 65.0| |
| |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]| 93.0| |
| |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|121.0| |
| +--------------------+------------------+-----+ |
| only showing top 5 rows |
| |
| * Multiple outputs |
| Some models can provide multiple outputs. These can be returned as a dictionary of named |
| values, which can be represented in either columnar or row-based formats. |
| |
| >>> def make_multi_sum_fn(): |
| ... def predict_columnar(x1: np.ndarray, x2: np.ndarray) -> Mapping[str, np.ndarray]: |
| ... # x1.shape = [batch_size, 4] |
| ... # x2.shape = [batch_size, 3] |
| ... return { |
| ... "sum1": np.sum(x1, axis=1), |
| ... "sum2": np.sum(x2, axis=1) |
| ... } |
| ... return predict_columnar |
| ... |
| >>> multi_sum_udf = predict_batch_udf( |
| ... make_multi_sum_fn, |
| ... return_type=StructType([ |
| ... StructField("sum1", FloatType(), True), |
| ... StructField("sum2", FloatType(), True) |
| ... ]), |
| ... batch_size=5, |
| ... input_tensor_shapes=[[4], [3]], |
| ... ) |
| >>> df.withColumn("preds", multi_sum_udf("t1", "t2")).select("t1", "t2", "preds.*").show(5) |
| +--------------------+------------------+----+----+ |
| | t1| t2|sum1|sum2| |
| +--------------------+------------------+----+----+ |
| |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 6.0| 3.0| |
| |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|22.0|15.0| |
| |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|38.0|27.0| |
| |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|54.0|39.0| |
| |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|70.0|51.0| |
| +--------------------+------------------+----+----+ |
| only showing top 5 rows |
| |
| >>> def make_multi_sum_fn(): |
| ... def predict_row(x1: np.ndarray, x2: np.ndarray) -> list[Mapping[str, float]]: |
| ... # x1.shape = [batch_size, 4] |
| ... # x2.shape = [batch_size, 3] |
| ... return [{'sum1': np.sum(x1[i]), 'sum2': np.sum(x2[i])} for i in range(len(x1))] |
| ... return predict_row |
| ... |
| >>> multi_sum_udf = predict_batch_udf( |
| ... make_multi_sum_fn, |
| ... return_type=StructType([ |
| ... StructField("sum1", FloatType(), True), |
| ... StructField("sum2", FloatType(), True) |
| ... ]), |
| ... batch_size=5, |
| ... input_tensor_shapes=[[4], [3]], |
| ... ) |
| >>> df.withColumn("sum", multi_sum_udf("t1", "t2")).select("t1", "t2", "sum.*").show(5) |
| +--------------------+------------------+----+----+ |
| | t1| t2|sum1|sum2| |
| +--------------------+------------------+----+----+ |
| |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 6.0| 3.0| |
| |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|22.0|15.0| |
| |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|38.0|27.0| |
| |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|54.0|39.0| |
| |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|70.0|51.0| |
| +--------------------+------------------+----+----+ |
| only showing top 5 rows |
| |
| Note that the multiple outputs can be arrays as well. |
| |
| >>> def make_multi_times_two_fn(): |
| ... def predict(x1: np.ndarray, x2: np.ndarray) -> Mapping[str, np.ndarray]: |
| ... # x1.shape = [batch_size, 4] |
| ... # x2.shape = [batch_size, 3] |
| ... return {"t1x2": x1 * 2, "t2x2": x2 * 2} |
| ... return predict |
| ... |
| >>> multi_times_two_udf = predict_batch_udf( |
| ... make_multi_times_two_fn, |
| ... return_type=StructType([ |
| ... StructField("t1x2", ArrayType(FloatType()), True), |
| ... StructField("t2x2", ArrayType(FloatType()), True) |
| ... ]), |
| ... batch_size=5, |
| ... input_tensor_shapes=[[4], [3]], |
| ... ) |
| >>> df.withColumn("x2", multi_times_two_udf("t1", "t2")).select("t1", "t2", "x2.*").show(5) |
| +--------------------+------------------+--------------------+------------------+ |
| | t1| t2| t1x2| t2x2| |
| +--------------------+------------------+--------------------+------------------+ |
| |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]|[0.0, 2.0, 4.0, 6.0]| [0.0, 2.0, 4.0]| |
| |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|[8.0, 10.0, 12.0,...| [8.0, 10.0, 12.0]| |
| |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|[16.0, 18.0, 20.0...|[16.0, 18.0, 20.0]| |
| |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|[24.0, 26.0, 28.0...|[24.0, 26.0, 28.0]| |
| |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|[32.0, 34.0, 36.0...|[32.0, 34.0, 36.0]| |
| +--------------------+------------------+--------------------+------------------+ |
| only showing top 5 rows |
| """ |
| # generate a new uuid each time this is invoked on the driver to invalidate executor-side cache. |
| model_uuid = uuid.uuid4() |
| |
| def predict(data: Iterator[Union[pd.Series, pd.DataFrame]]) -> Iterator[pd.DataFrame]: |
| # TODO: adjust return type hint when Iterator[Union[pd.Series, pd.DataFrame]] is supported |
| from pyspark.ml.model_cache import ModelCache |
| |
| # get predict function (from cache or from running user-provided make_predict_fn) |
| predict_fn = ModelCache.get(model_uuid) |
| if not predict_fn: |
| predict_fn = make_predict_fn() |
| ModelCache.add(model_uuid, predict_fn) |
| |
| # get number of expected parameters for predict function |
| signature = inspect.signature(predict_fn) |
| num_expected_cols = len(signature.parameters) |
| |
| # convert sparse input_tensor_shapes to dense if needed |
| input_shapes: List[List[int] | None] |
| if isinstance(input_tensor_shapes, Mapping): |
| input_shapes = [None] * num_expected_cols |
| for index, shape in input_tensor_shapes.items(): |
| input_shapes[index] = shape |
| else: |
| input_shapes = input_tensor_shapes # type: ignore |
| |
| # iterate over pandas batch, invoking predict_fn with ndarrays |
| for pandas_batch in data: |
| has_tuple = isinstance(pandas_batch, Tuple) # type: ignore |
| has_tensors = _has_tensor_cols(pandas_batch) |
| |
| # require input_tensor_shapes for any tensor columns |
| if has_tensors and not input_shapes: |
| raise ValueError("Tensor columns require input_tensor_shapes") |
| |
| for batch in _batched(pandas_batch, batch_size): |
| num_input_rows = len(batch) |
| num_input_cols = len(batch.columns) |
| if num_input_cols == num_expected_cols and num_expected_cols > 1: |
| # input column per expected input for multiple inputs |
| multi_inputs = _validate_and_transform_multiple_inputs( |
| batch, input_shapes, num_input_cols |
| ) |
| # run model prediction function on multiple (numpy) inputs |
| preds = predict_fn(*multi_inputs) |
| elif num_expected_cols == 1: |
| # one or more input columns for single expected input |
| single_input = _validate_and_transform_single_input( |
| batch, input_shapes, has_tensors, has_tuple |
| ) |
| # run model prediction function on single (numpy) inputs |
| preds = predict_fn(single_input) |
| else: |
| msg = "Model expected {} inputs, but received {} columns" |
| raise ValueError(msg.format(num_expected_cols, num_input_cols)) |
| |
| # return transformed predictions to Spark |
| yield _validate_and_transform_prediction_result( |
| preds, num_input_rows, return_type |
| ) # type: ignore |
| |
| return pandas_udf(predict, return_type) # type: ignore[call-overload] |
| |
| |
| def _test() -> None: |
| import doctest |
| from pyspark.sql import SparkSession |
| import pyspark.ml.functions |
| import sys |
| |
| from pyspark.sql.pandas.utils import ( |
| require_minimum_pandas_version, |
| require_minimum_pyarrow_version, |
| ) |
| |
| try: |
| require_minimum_pandas_version() |
| require_minimum_pyarrow_version() |
| except Exception as e: |
| print( |
| f"Skipping pyspark.ml.functions doctests: {e}", |
| file=sys.stderr, |
| ) |
| sys.exit(0) |
| |
| globs = pyspark.ml.functions.__dict__.copy() |
| spark = SparkSession.builder.master("local[2]").appName("ml.functions tests").getOrCreate() |
| sc = spark.sparkContext |
| globs["sc"] = sc |
| globs["spark"] = spark |
| |
| (failure_count, test_count) = doctest.testmod( |
| pyspark.ml.functions, |
| globs=globs, |
| optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, |
| ) |
| spark.stop() |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
