I would like to implement a model based on some cleaned and prepared data set. I already have a bit of experience with PySpark, but from a data scientist's perspective it can be cumbersome to work with it. Therefore I would like to try Koalas. I have a lot of experience with Pandas and hope this API will help me to leverage my skills. My question is: does the Koalas library allow to use all Pandas machine learning libraries like Scikit-Learn, XGBoost, and TensorFlow ? As the Koalas API is relatively new, what draw-backs can be expected besides the lack of documentation ?


2 Answers 2


Scikit-Learn, XGBoost and TensorFlow don't work with Koalas DataFrames directly. But you can use them with MlFlow. Here is an example of ML model where inference was done with Koalas:

from mlflow.tracking import MlflowClient, set_tracking_uri
import mlflow.sklearn
from tempfile import mkdtemp
d = mkdtemp("koalas_mlflow")
client = MlflowClient()
exp = mlflow.create_experiment("my_experiment")

from sklearn.linear_model import LinearRegression
train = pd.DataFrame({"x1": np.arange(8), "x2": np.arange(8)**2,
                      "y": np.log(2 + np.arange(8))})
train_x = train[["x1", "x2"]]
train_y = train[["y"]]
with mlflow.start_run():
    lr = LinearRegression()
    lr.fit(train_x, train_y)
    mlflow.sklearn.log_model(lr, "model")

from databricks.koalas.mlflow import load_model
run_info = client.list_run_infos(exp)[-1]
model = load_model("runs:/{run_id}/model".format(run_id=run_info.run_uuid))
prediction_df = ks.DataFrame({"x1": [2.0], "x2": [4.0]})
prediction_df["prediction"] = model.predict(prediction_df)

Probably you can do the same with TensorFlow as well.

Another option is to use Spark ML with Koalas by converting Koalas dataframes into Spark dataframes.

Or you can use Pandas UDFs instead of Spark ML. In this case you can incorporate Koalas with Tensorflow.


I think you have misunderstood the koalas library. You can say its Pandas on Distributed System. You can use Koalas similar to pandas. There are few drawbacks with respect to APIs which is documented in their docs and few articles already written on medium.

You can do your EDA and straight away use them in all the libraries you have mentioned.

Recent pandas 1.0 is faster compared to older versions. It also uses Numba Behind the scene.

Vaex is another library which is available and you can use for EDA but the api names are different from pandas.

Also you can use Modin and dask and use it like pandas with few limitations again.

Clearly these libraries have no dependencies on Sklearn, XGBOOST or TF. You can split train and run your model.

  • $\begingroup$ Thank you Syenix, that helped me clarify things. I will look for articles and see how it "works under the hood". $\endgroup$
    – DataBach
    Feb 15, 2020 at 17:59
  • $\begingroup$ When you are satisfied, could you please accept it as answer :) $\endgroup$
    – Syenix
    Feb 16, 2020 at 13:31

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.