code: def get_prediction_udf(
      model, method="predict", feature_type="numpy", names=None
  ):
      """
      Build a vectorized PySpark UDF to apply a sklearn model's `predict` or
      `predict_proba` methods columns in a PySpark DataFrame. Handles
      flexible types of feature data for prediction including 2-D numpy
      arrays ('numpy'), single field text data ('text') and mixed type
      pandas DataFrames ('pandas'). The UDF can then be applied as shown in the
      example below.
  
      NOTE: This function requires pyarrow and pyspark with appropriate
      versions for vectorized pandas UDFs and appropriate spark configuration
      to use pyarrow. Ths requires pyarrow>=0.8.0 and pyspark>=2.3.0.
      Additionally, the spark version must be 2.3 or higher. These requirements
      are not enforced by the sk-dist package at setup time.
  
      Args:
          model (sklearn Estimator): sklearn model to distribute
              predictions with PySpark
          method (str): name of prediction method; either 'predict'
              or 'predict_proba'
          feature_type (str): name of feature type; either 'numpy',
              'pandas' or 'text'
          names (array-like): list of ordered column names
              (only necessary for 'pandas' feature_type
      Returns:
          PySpark pandas UDF (pyspark.sql.functions.pandas_udf)
      Example:
      >>> import pandas as pd
      >>> from sklearn.datasets import load_digits
      >>> from sklearn.linear_model import LogisticRegression
      >>> from pyspark.sql import SparkSession
      >>> spark = (
      >>>     SparkSession
      >>>     .builder
      >>>     .getOrCreate()
      >>>     )
      >>> data = load_digits()
      >>> X = data["data"]
      >>> y = data["target"]
      >>> model = LogisticRegression()
      >>> model.fit(X, y)
      >>> predict = get_prediction_udf(model, method="predict")
      >>> predict_proba = get_prediction_udf(model, method="predict_proba")
      >>> pdf = pd.DataFrame(X)
      >>> sdf = spark.createDataFrame(pdf)
      >>> cols = [F.col(str(c)) for c in sdf.columns]
      >>> prediction_df = (
      >>>     sdf
      >>>     .withColumn("scores", predict_proba(*cols))
      >>>     .withColumn("preds", predict(*cols))
      >>>     .select("preds", "scores")
      >>>     )
      >>> prediction_df.show()
      ... +-----+--------------------+
      ... |preds|              scores|
      ... +-----+--------------------+
      ... |    0|[0.99988026795692...|
      ... |    1|[4.75035277837040...|
      ... |    2|[2.94811218592164...|
      ... |    3|[1.63438595023762...|
      ... |    4|[1.11339868338047...|
      ... |    5|[1.47300432716012...|
      ... |    6|[1.08560009259480...|
      ... |    7|[3.02428232165044...|
      ... |    8|[7.65445972596079...|
      ... |    9|[3.97610488897298...|
      ... |    0|[0.99918670844137...|
      ... |    1|[2.65336456879078...|
      ... |    2|[1.85886361541580...|
      ... |    3|[2.89824009324990...|
      ... |    4|[2.84813979824305...|
      ... |    5|[2.70090567992820...|
      ... |    6|[1.10907772018062...|
      ... |    7|[3.06455862370095...|
      ... |    8|[2.38739344440480...|
      ... |    9|[8.23628591704589...|
      ... +-----+--------------------+
      ... only showing top 20 rows
      """
      if not _is_pyspark_installed():
          raise ImportError("Module pyspark not found")
      if not _is_pyarrow_installed():
          raise ImportError("Module pyarrow not found")
      from pyspark.sql import functions as F
      from pyspark.sql.types import (
          DoubleType,
          StringType,
          IntegerType,
          ArrayType,
      )
  
      if method == "predict":
  
          def predict_func(*cols):
              vals = _get_vals(
                  *cols, feature_type=feature_type, names=names
              )
              return pd.Series(model.predict(vals))
  
          return_type = (
              StringType()
              if isinstance(model.classes_[0], str)
              else IntegerType()
          )
          predict = F.pandas_udf(predict_func, returnType=return_type)
      elif method == "predict_proba":
  
          def predict_func(*cols):
              vals = _get_vals(
                  *cols, feature_type=feature_type, names=names
              )
              return pd.Series(list(model.predict_proba(vals)))
  
          predict = F.pandas_udf(
              predict_func, returnType=ArrayType(DoubleType())
          )
      else:
          raise ValueError("Unknown method: {0}".format(method))
      return predict
  
file_path: sk-dist/skdist/distribute/predict.py
imports: import pandas
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: get_prediction_udf(model, method, feature_type, names) and returns: predict, pd, pd