Set Dask-sql as the default Fugue Dask engine when installed

continuous_integration/environment-3.10-jdk11-dev.yaml

@@ -31,14 +31,5 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0
+  - fugue>=0.7.0

continuous_integration/environment-3.10-jdk8-dev.yaml

@@ -31,14 +31,5 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0
+  - fugue>=0.7.0

continuous_integration/environment-3.8-jdk11-dev.yaml

@@ -31,13 +31,5 @@ dependencies:
 - tpot
 - tzlocal=2.1
 - uvicorn=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  - fugue[sql]==0.5.3
+  - fugue>=0.7.0

continuous_integration/environment-3.8-jdk8-dev.yaml

@@ -31,13 +31,5 @@ dependencies:
 - tpot
 - tzlocal=2.1
 - uvicorn=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  - fugue[sql]==0.5.3
+  - fugue>=0.7.0

continuous_integration/environment-3.9-jdk11-dev.yaml

@@ -31,14 +31,5 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0
+  - fugue>=0.7.0

continuous_integration/environment-3.9-jdk8-dev.yaml

@@ -31,14 +31,5 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
 - pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0
+  - fugue>=0.7.0

dask_sql/integrations/fugue.py

@@ -1,8 +1,10 @@
 try:
     import fugue
     import fugue_dask
+    from dask.distributed import Client
     from fugue import WorkflowDataFrame, register_execution_engine
     from fugue_sql import FugueSQLWorkflow
+    from triad import run_at_def
     from triad.utils.convert import get_caller_global_local_vars
 except ImportError:  # pragma: no cover
     raise ImportError(
@@ -15,9 +17,24 @@
 
 from dask_sql.context import Context
 
-register_execution_engine(
-    "dask", lambda conf: DaskSQLExecutionEngine(conf), on_dup="overwrite"
-)
+
+@run_at_def
+def register_engines() -> None:
+    """Register (overwrite) the default Dask execution engine of Fugue. This
+    function is invoked as an entrypoint, users don't need to call it explicitly.
+    """
+    register_execution_engine(
+        "dask",
+        lambda conf, **kwargs: DaskSQLExecutionEngine(conf=conf),
+        on_dup="overwrite",
+    )
+    register_execution_engine(
+        Client,
+        lambda engine, conf, **kwargs: DaskSQLExecutionEngine(
+            dask_client=engine, conf=conf
+        ),
+        on_dup="overwrite",
+    )
 
 
 class DaskSQLEngine(fugue.execution.execution_engine.SQLEngine):

docs/environment.yml

@@ -10,8 +10,7 @@ dependencies:
   - maven>=3.6.0
   - dask>=2022.3.0
   - pandas>=1.1.2
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue>=0.5.3,<0.7.0
+  - fugue>=0.7.0
   - jpype1>=1.0.2
   - fastapi>=0.69.0
   - uvicorn>=0.11.3

docs/requirements-docs.txt

@@ -3,8 +3,7 @@ sphinx-tabs
 dask-sphinx-theme>=3.0.0
 dask>=2022.3.0
 pandas>=1.1.2
-# FIXME: tests are failing with fugue 0.7.0
-fugue>=0.5.3,<0.7.0
+fugue>=0.7.0
 jpype1>=1.0.2
 fastapi>=0.69.0
 uvicorn>=0.11.3

notebooks/Feature Overview.ipynb

@@ -748,9 +748,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "Python 3"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -762,7 +762,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.6"
+   "version": "3.10.4"
   },
   "toc-autonumbering": false,
   "toc-showcode": false,