Machine Learning utilities


Arbitrary MLflow models can be used with Koalas Dataframes, provided they implement the ‘pyfunc’ flavor. This is the case for most frameworks supported by MLflow (scikit-learn, pytorch, tensorflow, …). See comprehensive examples in load_model() for more information.


The MLflow package must be installed in order to use this module. If MLflow is not installed in your environment already, you can install it with the following command:

pip install koalas[mlflow]

PythonModelWrapper(model_uri, return_type_hint)

A wrapper around MLflow’s Python object model.

load_model(model_uri[, predict_type])

Loads an MLflow model into an wrapper that can be used both for pandas and Koalas DataFrame.