Does the Koalas library allow to use all Pandas machine learning libraries like Scikit-Learn, XGBoost, and TensorFlow?

Question

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 ?

Answer 1

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")
set_tracking_uri("file:%s"%d)
client = MlflowClient()
exp = mlflow.create_experiment("my_experiment")
mlflow.set_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)
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.

Answer 2

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.

Answer 3

I personally found the accepted solution unnecessary complicated to understand. I don't know why the user used MLFlow. The following code achieves the same thing in a more understandable way:

import numpy as np
import pandas as pd
import koalas as ks
from sklearn.linear_model import LinearRegression

# Create a pandas dataframe
train = pd.DataFrame({
    "x1": np.arange(8), "x2": np.arange(8)**2,
    "y": np.log(2 + np.arange(8))
})

# Convert the pandas dataframe to a Koalas dataframe
train_ks = ks.from_pandas(train)

# Split the Koalas dataframe into input and output data
train_x = train_ks[["x1", "x2"]]
train_y = train_ks[["y"]]

# Train a linear regression model with scikit-learn
lr = LinearRegression()
lr.fit(train_x, train_y)

# Use the trained model to make predictions on a new data set
input_data = ks.DataFrame({"x1": [2.0], "x2": [4.0]})
input_data["prediction"] = lr.predict(input_data)
print(input_data)

The predict method of a scikit-learn model can be used on a Koalas dataframe, as long as the Koalas dataframe has the same structure as the input data used to train the model.

The main drawbacks with Koalas are that:

1. It aims to provide a Pandas-like experience, but may not have the same performance as PySpark in certain situations, especially when dealing with large data sets or complex operations
2. Not all PySpark functionality are available in Koalas