code: def test_predict(spark_session):
      sc = spark_session.sparkContext
  
      # simple 2-D numpy features
      data = load_digits()
      X = data["data"]
      y = data["target"]
      model = LogisticRegression(solver="liblinear", multi_class="auto")
      model.fit(X, y)
  
      # get UDFs with default 'numpy' feature types
      predict = get_prediction_udf(model, method="predict")
      predict_proba = get_prediction_udf(model, method="predict_proba")
  
      # create PySpark DataFrame from features
      pdf = pd.DataFrame(X)
      sdf = spark_session.createDataFrame(pdf)
      cols = [F.col(str(c)) for c in sdf.columns]
  
      # apply predict UDFs and select prediction output
      prediction_df = (
          sdf.withColumn("scores", predict_proba(*cols))
          .withColumn("preds", predict(*cols))
          .select("preds", "scores")
      )
      assert prediction_df.count() == X.shape[0]
  
file_path: sk-dist/skdist/tests/test_spark.py
imports: import pytest
  import sys
  import pandas
  import numpy
  from sklearn.datasets import load_breast_cancer, load_digits
  from sklearn.model_selection import train_test_split
  from sklearn.linear_model import LogisticRegression
  from skdist.distribute.multiclass import DistOneVsRestClassifier
  from skdist.distribute.search import DistGridSearchCV, DistRandomizedSearchCV
  from skdist.distribute.ensemble import DistRandomForestClassifier
  from skdist.distribute.predict import get_prediction_udf
label_desc: Use image node_id 6 for calling a global function with example usage: test_predict(spark_session) without return types