Multiple schemas allowed (#205)

* ML model improvement : Adding "SHOW MODELS and DESCRIBE MODEL"

Author:    rajagurunath <[email protected]>
Date:      Mon May 24 02:37:40 2021 +0530

* fix typo

* ML model improvement : added EXPORT MODEL

* ML model improvement : refactoring for PR

* ML model improvement : Adding stmts in notebook

* ML model improvement : Adding stmts in notebook

* ML model improvement : also test the non-happy path

* ML model improvement : Added mlflow and <With> in sql for extra params

* ML model improvement : Added mlflow and <With> in sql for extra params

* Added Test cases for Export MODEL

* Added ML documentation about the following:
1. SHOW MODELS
2. DESCRIBE MODEL
3. EXPORT MODEL

* refactored based on PR

* Added support only for sklearn compatible models

* excluded mlflow part from code coverage

* install mlflow in test cluster

* Added test for non sklearn compatible model

* Added: initial draft of referencing multiple schemas

* Added schema DDLs
1. Create Schema
2. Use schema
3. Drop schema
4. Added testcases

* Use compound identifier for models, tables, experiments, views

* Split the compound identifiers - without using the schema so far

* Added a schema_name parameter to most functions and actually use the schema

* Pass on the schemas to JAVA

* Some simplifications and tests

* Some cleanup, documentation and more tests (and fixed a bug in
aggregation)

* Remove unneeded import

Co-authored-by: gurunath <[email protected]>
Co-authored-by: Gurunath LankupalliVenugopal <[email protected]>

Author: 3 people

dask_sql/context.py

@@ -1,9 +1,15 @@
-from typing import Dict, List, Tuple, Union
+from collections import namedtuple
+from typing import Any, Callable, Dict, List, Tuple, Union
 
 import dask.dataframe as dd
+import pandas as pd
 
 ColumnType = Union[str, int]
 
+FunctionDescription = namedtuple(
+    "FunctionDescription", ["name", "parameters", "return_type", "aggregation"]
+)
+
 
 class ColumnContainer:
     # Forward declaration
@@ -173,3 +179,13 @@ def assign(self) -> dd.DataFrame:
             }
         )
         return df[self.column_container.columns]
+
+
+class SchemaContainer:
+    def __init__(self, name: str):
+        self.__name__ = name
+        self.tables: Dict[str, DataContainer] = {}
+        self.experiments: Dict[str, pd.DataFrame] = {}
+        self.models: Dict[str, Tuple[Any, List[str]]] = {}
+        self.functions: Dict[str, Callable] = {}
+        self.function_lists: List[FunctionDescription] = []

dask_sql/datacontainer.py

@@ -113,7 +113,11 @@ def median(df:pd.DataFrame) -> pd.DataFrame:
     _global, _local = get_caller_global_local_vars()
 
     dag = FugueSQLWorkflow()
-    dfs = {} if ctx is None else {k: dag.df(v.df) for k, v in ctx.tables.items()}
+    dfs = (
+        {}
+        if ctx is None
+        else {k: dag.df(v.df) for k, v in ctx.schema[ctx.schema_name].tables.items()}
+    )
     result = dag._sql(sql, _global, _local, **dfs)
     dag.run(DaskSQLExecutionEngine(conf=fugue_conf))
 

dask_sql/integrations/fugue.py

@@ -90,6 +90,9 @@ def _set_or_check_java_home():
 DaskScalarFunction = com.dask.sql.schema.DaskScalarFunction
 DaskSchema = com.dask.sql.schema.DaskSchema
 RelationalAlgebraGenerator = com.dask.sql.application.RelationalAlgebraGenerator
+RelationalAlgebraGeneratorBuilder = (
+    com.dask.sql.application.RelationalAlgebraGeneratorBuilder
+)
 SqlTypeName = org.apache.calcite.sql.type.SqlTypeName
 ValidationException = org.apache.calcite.tools.ValidationException
 SqlParseException = org.apache.calcite.sql.parser.SqlParseException

dask_sql/java.py

@@ -2,30 +2,36 @@
 from .columns import ShowColumnsPlugin
 from .create_experiment import CreateExperimentPlugin
 from .create_model import CreateModelPlugin
+from .create_schema import CreateSchemaPlugin
 from .create_table import CreateTablePlugin
 from .create_table_as import CreateTableAsPlugin
 from .describe_model import ShowModelParamsPlugin
 from .drop_model import DropModelPlugin
+from .drop_schema import DropSchemaPlugin
 from .drop_table import DropTablePlugin
 from .export_model import ExportModelPlugin
 from .predict import PredictModelPlugin
 from .schemas import ShowSchemasPlugin
 from .show_models import ShowModelsPlugin
+from .switch_schema import SwitchSchemaPlugin
 from .tables import ShowTablesPlugin
 
 __all__ = [
     AnalyzeTablePlugin,
+    CreateExperimentPlugin,
     CreateModelPlugin,
+    CreateSchemaPlugin,
     CreateTableAsPlugin,
     CreateTablePlugin,
     DropModelPlugin,
+    DropSchemaPlugin,
     DropTablePlugin,
+    ExportModelPlugin,
     PredictModelPlugin,
     ShowColumnsPlugin,
+    ShowModelParamsPlugin,
+    ShowModelsPlugin,
     ShowSchemasPlugin,
     ShowTablesPlugin,
-    ShowModelsPlugin,
-    ShowModelParamsPlugin,
-    ExportModelPlugin,
-    CreateExperimentPlugin,
+    SwitchSchemaPlugin,
 ]

dask_sql/physical/rel/custom/__init__.py

@@ -4,7 +4,6 @@
 from dask_sql.datacontainer import ColumnContainer, DataContainer
 from dask_sql.mappings import python_to_sql_type
 from dask_sql.physical.rel.base import BaseRelPlugin
-from dask_sql.utils import get_table_from_compound_identifier
 
 
 class AnalyzeTablePlugin(BaseRelPlugin):
@@ -28,8 +27,8 @@ class AnalyzeTablePlugin(BaseRelPlugin):
     def convert(
         self, sql: "org.apache.calcite.sql.SqlNode", context: "dask_sql.Context"
     ) -> DataContainer:
-        components = list(map(str, sql.getTableName().names))
-        dc = get_table_from_compound_identifier(context, components)
+        schema_name, name = context.fqn(sql.getTableName())
+        dc = context.schema[schema_name].tables[name]
         columns = list(map(str, sql.getColumnList()))
 
         if not columns:

dask_sql/physical/rel/custom/analyze.py

@@ -4,7 +4,6 @@
 from dask_sql.datacontainer import ColumnContainer, DataContainer
 from dask_sql.mappings import python_to_sql_type
 from dask_sql.physical.rel.base import BaseRelPlugin
-from dask_sql.utils import get_table_from_compound_identifier
 
 
 class ShowColumnsPlugin(BaseRelPlugin):
@@ -22,8 +21,8 @@ class ShowColumnsPlugin(BaseRelPlugin):
     def convert(
         self, sql: "org.apache.calcite.sql.SqlNode", context: "dask_sql.Context"
     ) -> DataContainer:
-        components = list(map(str, sql.getTable().names))
-        dc = get_table_from_compound_identifier(context, components)
+        schema_name, name = context.fqn(sql.getTable())
+        dc = context.schema[schema_name].tables[name]
 
         cols = dc.column_container.columns
         dtypes = list(map(lambda x: str(python_to_sql_type(x)).lower(), dc.df.dtypes))

dask_sql/physical/rel/custom/columns.py

@@ -98,10 +98,10 @@ def convert(
         self, sql: "org.apache.calcite.sql.SqlNode", context: "dask_sql.Context"
     ) -> DataContainer:
         select = sql.getSelect()
-        experiment_name = str(sql.getExperimentName())
+        schema_name, experiment_name = context.fqn(sql.getExperimentName())
         kwargs = convert_sql_kwargs(sql.getKwargs())
 
-        if experiment_name in context.experiments:
+        if experiment_name in context.schema[schema_name].experiments:
             if sql.getIfNotExists():
                 return
             elif not sql.getReplace():
@@ -175,6 +175,7 @@ def convert(
                 experiment_name,
                 ParallelPostFit(estimator=search.best_estimator_),
                 X.columns,
+                schema_name=schema_name,
             )
 
         if automl_class:
@@ -198,9 +199,12 @@ def convert(
                 experiment_name,
                 ParallelPostFit(estimator=automl.fitted_pipeline_),
                 X.columns,
+                schema_name=schema_name,
             )
 
-        context.register_experiment(experiment_name, experiment_results=df)
+        context.register_experiment(
+            experiment_name, experiment_results=df, schema_name=schema_name
+        )
         cc = ColumnContainer(df.columns)
         dc = DataContainer(dd.from_pandas(df, npartitions=1), cc)
         return dc

dask_sql/physical/rel/custom/create_experiment.py

@@ -105,10 +105,10 @@ def convert(
         self, sql: "org.apache.calcite.sql.SqlNode", context: "dask_sql.Context"
     ) -> DataContainer:
         select = sql.getSelect()
-        model_name = str(sql.getModelName())
+        schema_name, model_name = context.fqn(sql.getModelName())
         kwargs = convert_sql_kwargs(sql.getKwargs())
 
-        if model_name in context.models:
+        if model_name in context.schema[schema_name].models:
             if sql.getIfNotExists():
                 return
             elif not sql.getReplace():
@@ -162,4 +162,4 @@ def convert(
             y = None
 
         model.fit(X, y, **fit_kwargs)
-        context.register_model(model_name, model, X.columns)
+        context.register_model(model_name, model, X.columns, schema_name=schema_name)

dask_sql/physical/rel/custom/create_model.py

@@ -0,0 +1,40 @@
+import logging
+
+from dask_sql.datacontainer import DataContainer
+from dask_sql.physical.rel.base import BaseRelPlugin
+
+logger = logging.getLogger(__name__)
+
+
+class CreateSchemaPlugin(BaseRelPlugin):
+    """
+    Create a schema with the given name
+    and register it at the context.
+    The SQL call looks like
+
+        CREATE SCHEMA <schema-name>
+
+    Using this SQL is equivalent to just doing
+
+        context.create_schema(<schema-name>)
+
+    but can also be used without writing a single line of code.
+    Nothing is returned.
+    """
+
+    class_name = "com.dask.sql.parser.SqlCreateSchema"
+
+    def convert(
+        self, sql: "org.apache.calcite.sql.SqlNode", context: "dask_sql.Context"
+    ):
+        schema_name = str(sql.getSchemaName())
+
+        if schema_name in context.schema:
+            if sql.getIfNotExists():
+                return
+            elif not sql.getReplace():
+                raise RuntimeError(
+                    f"A Schema with the name {schema_name} is already present."
+                )
+
+        context.create_schema(schema_name)