Backend agnostic machine learning models

dask_sql/physical/rel/custom/create_experiment.py

@@ -6,6 +6,7 @@
 
 from dask_sql.datacontainer import ColumnContainer, DataContainer
 from dask_sql.physical.rel.base import BaseRelPlugin
+from dask_sql.physical.rel.custom.ml_classes import get_cpu_classes, get_gpu_classes
 from dask_sql.utils import convert_sql_kwargs, import_class
 
 if TYPE_CHECKING:
@@ -14,6 +15,9 @@
 
 logger = logging.getLogger(__name__)
 
+cpu_classes = get_cpu_classes()
+gpu_classes = get_gpu_classes()
+
 
 class CreateExperimentPlugin(BaseRelPlugin):
     """
@@ -147,6 +151,13 @@ def convert(self, rel: "LogicalPlan", context: "dask_sql.Context") -> DataContai
         y = training_df[target_column]
 
         if model_class and experiment_class:
+            if type(training_df) == dd.core.DataFrame:
+                model_class = cpu_classes.get(model_class, model_class)
+                experiment_class = cpu_classes.get(experiment_class, experiment_class)
+            elif "cudf" in str(training_df._partition_type):
+                model_class = gpu_classes.get(model_class, model_class)
+                experiment_class = gpu_classes.get(experiment_class, experiment_class)
+
             try:
                 ModelClass = import_class(model_class)
             except ImportError:

dask_sql/physical/rel/custom/create_model.py

@@ -1,11 +1,13 @@
 import logging
 from typing import TYPE_CHECKING
 
+import dask.dataframe as dd
 import numpy as np
 from dask import delayed
 
 from dask_sql.datacontainer import DataContainer
 from dask_sql.physical.rel.base import BaseRelPlugin
+from dask_sql.physical.rel.custom.ml_classes import get_cpu_classes, get_gpu_classes
 from dask_sql.utils import convert_sql_kwargs, import_class
 
 if TYPE_CHECKING:
@@ -14,6 +16,9 @@
 
 logger = logging.getLogger(__name__)
 
+cpu_classes = get_cpu_classes()
+gpu_classes = get_gpu_classes()
+
 
 class CreateModelPlugin(BaseRelPlugin):
     """
@@ -141,6 +146,11 @@ def convert(self, rel: "LogicalPlan", context: "dask_sql.Context") -> DataContai
             X = training_df
             y = None
 
+        if type(training_df) == dd.core.DataFrame:
+            model_class = cpu_classes.get(model_class, model_class)
+        elif "cudf" in str(training_df._partition_type):
+            model_class = gpu_classes.get(model_class, model_class)
+
         try:
             ModelClass = import_class(model_class)
         except ImportError: