Function create_spark_df

demo/guide-python/spark_estimator_examples.py:21–27

(X: np.ndarray, y: np.ndarray)

def create_spark_df(X: np.ndarray, y: np.ndarray) -> DataFrame:
    return spark.createDataFrame(
        spark.sparkContext.parallelize(
            [(Vectors.dense(features), float(label)) for features, label in zip(X, y)]
        ),
        ["features", "label"],
    )


# load diabetes dataset (regression dataset)